Update app.py

app.py

@@ -19,111 +19,84 @@ os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 os.makedirs(RESULT_FOLDER, exist_ok=True)
 
 def upload_image(image_data_url):
-    """Handle base64 image upload and Gemini file upload"""
+    """Helper function to upload image to Gemini"""
     try:
         header, encoded = image_data_url.split(',', 1)
-        binary_data = base64.b64decode(encoded)
-        ext = ".png" if "png" in header.lower() else ".jpg"
-        temp_filename = secure_filename(f"temp_{os.urandom(8).hex()}{ext}")
-        temp_filepath = os.path.join(UPLOAD_FOLDER, temp_filename)
-        
-        with open(temp_filepath, "wb") as f:
-            f.write(binary_data)
-        
-        uploaded_file = client.files.upload(file=temp_filepath)
-        os.remove(temp_filepath)  # Clean up temporary file
-        return uploaded_file
+    except ValueError:
+        raise ValueError("Invalid image data")
     
-    except Exception as e:
-        raise ValueError(f"Image processing error: {str(e)}")
+    binary_data = base64.b64decode(encoded)
+    ext = ".png" if "png" in header.lower() else ".jpg"
+    temp_filename = secure_filename("temp_image" + ext)
+    temp_filepath = os.path.join(UPLOAD_FOLDER, temp_filename)
+    
+    with open(temp_filepath, "wb") as f:
+        f.write(binary_data)
+    
+    return client.files.upload(file=temp_filepath)
 
 def is_prohibited_request(uploaded_file, object_type):
-    """Check if request involves people/animals or their belongings"""
-    model = "gemini-2.0-flash-lite"
-    parts = [
-        types.Part.from_uri(file_uri=uploaded_file.uri, mime_type=uploaded_file.mime_type),
-        types.Part.from_text(text=f"Remove {object_type}")
+    """Check if request matches prohibited removal cases"""
+    object_type = object_type.lower()
+    
+    # Prohibited cases
+    prohibited_requests = [
+        "remove sunglasses" in object_type and "table" not in object_type,  # ❌ when worn
+        "remove phone" in object_type and "hand" in object_type,           # ❌ when in hand
+        "remove eyes" in object_type                                      # ❌ remove eyes
     ]
     
-    contents = [types.Content(role="user", parts=parts)]
+    # Allowed cases
+    allowed_requests = [
+        "remove sunglasses" in object_type and "table" in object_type,     # ✅ when on table
+        "remove car" in object_type,                                      # ✅ remove car
+        "remove background" in object_type                                # ✅ remove background
+    ]
     
-    generate_content_config = types.GenerateContentConfig(
-        system_instruction=[types.Part.from_text(text="""Analyze image and request to detect:
-1. Direct removal of people/animals
-2. Removal of items attached to/worn by people/animals
-3. Removal of body parts or personal belongings
-
-Prohibited examples:
-- Person, dog, cat
-- Sunglasses on face, mask, hat
-- Phone in hand, watch on wrist
-- Eyes, hands, hair
-
-Allowed examples:
-- Background, car, tree
-- Sunglasses on table
-- Phone on desk
-
-Respond ONLY with 'Yes' or 'No'""")],
-        temperature=0.0,
-        max_output_tokens=1,
-    )
+    # Check for person/animal removal
+    person_animal_check = "remove person" in object_type or "remove animal" in object_type or \
+                         "remove dog" in object_type or "remove cat" in object_type
     
-    try:
-        response = client.models.generate_content(
-            model=model,
-            contents=contents,
-            config=generate_content_config
-        )
-        if response.candidates and response.candidates[0].content.parts:
-            return response.candidates[0].content.parts[0].text.strip().lower() == "yes"
-        return True  # Default to safe mode if uncertain
-    except Exception as e:
-        print(f"Safety check failed: {str(e)}")
-        return True  # Block if check fails
+    return any(prohibited_requests) or person_animal_check
 
-def generate_modified_image(uploaded_file, object_type):
-    """Generate image with object removed using experimental model"""
+def generate_gemini_output(object_type, uploaded_file):
+    """Generate image using gemini-2.0-flash-exp-image-generation"""
     model = "gemini-2.0-flash-exp-image-generation"
     parts = [
         types.Part.from_uri(file_uri=uploaded_file.uri, mime_type=uploaded_file.mime_type),
-        types.Part.from_text(text=f"Completely remove {object_type} from the image without leaving traces")
+        types.Part.from_text(text=f"Remove {object_type} from the image")
     ]
     
     contents = [types.Content(role="user", parts=parts)]
     
     generate_content_config = types.GenerateContentConfig(
-        temperature=0.5,
-        top_p=0.9,
-        max_output_tokens=1024,
-        response_modalities=["image"],
+        temperature=1,
+        top_p=0.95,
+        top_k=40,
+        max_output_tokens=8192,
+        response_modalities=["image", "text"],
         safety_settings=[
-            types.SafetySetting(category="HARM_CATEGORY_CIVIC_INTEGRITY", threshold="BLOCK_NONE"),
-            types.SafetySetting(category="HARM_CATEGORY_VIOLENCE", threshold="BLOCK_NONE")
-        ]
+            types.SafetySetting(category="HARM_CATEGORY_CIVIC_INTEGRITY", threshold="OFF"),
+        ],
     )
     
-    try:
-        for chunk in client.models.generate_content_stream(
-            model=model,
-            contents=contents,
-            config=generate_content_config,
-        ):
-            if chunk.candidates and chunk.candidates[0].content.parts:
-                part = chunk.candidates[0].content.parts[0]
-                if part.inline_data:
-                    ext = mimetypes.guess_extension(part.inline_data.mime_type) or ".png"
-                    output_filename = secure_filename(f"result_{os.urandom(4).hex()}{ext}")
-                    output_path = os.path.join(RESULT_FOLDER, output_filename)
-                    
-                    with open(output_path, "wb") as f:
-                        f.write(part.inline_data.data)
-                    
-                    return output_path
-        return None
-    except Exception as e:
-        print(f"Image generation failed: {str(e)}")
-        return None
+    result_image = None
+    for chunk in client.models.generate_content_stream(
+        model=model,
+        contents=contents,
+        config=generate_content_config,
+    ):
+        if chunk.candidates and chunk.candidates[0].content.parts:
+            part = chunk.candidates[0].content.parts[0]
+            if part.inline_data:
+                file_extension = mimetypes.guess_extension(part.inline_data.mime_type) or ".png"
+                output_filename = secure_filename("generated_output" + file_extension)
+                result_image_path = os.path.join(RESULT_FOLDER, output_filename)
+                with open(result_image_path, "wb") as f:
+                    f.write(part.inline_data.data)
+                result_image = result_image_path
+    
+    return result_image
 
 @app.route("/")
 def index():
@@ -132,40 +105,36 @@ def index():
 @app.route("/process", methods=["POST"])
 def process():
     try:
-        data = request.get_json()
-        if not data or "image" not in data or "objectType" not in data:
-            return jsonify({"success": False, "message": "Invalid request format"}), 400
-            
-        image_data = data["image"]
-        object_type = data["objectType"].strip().lower()
+        data = request.get_json(force=True)
+        image_data = data.get("image")
+        object_type = data.get("objectType", "").strip().lower()
         
-        if not object_type:
-            return jsonify({"success": False, "message": "Please specify an object to remove"}), 400
+        if not image_data or not object_type:
+            return jsonify({"success": False, "message": "Missing required data"}), 400
         
-        # Process image upload
+        # Upload image once
         uploaded_file = upload_image(image_data)
         
-        # Safety check
+        # Check for prohibited requests
         if is_prohibited_request(uploaded_file, object_type):
-            return jsonify({
-                "success": False,
-                "message": "Cannot remove people, animals, or personal items"
-            }), 403
+            error_message = "Sorry, I can't assist with this request."
+            if "person" in object_type or "animal" in object_type or "cat" in object_type or "dog" in object_type:
+                error_message = "Sorry, I can't assist with removing people or animals."
+            return jsonify({"success": False, "message": error_message}), 400
+        
+        # Generate output if allowed
+        result_image = generate_gemini_output(object_type, uploaded_file)
         
-        # Generate modified image
-        result_path = generate_modified_image(uploaded_file, object_type)
-        if not result_path:
+        if not result_image:
             return jsonify({"success": False, "message": "Failed to generate image"}), 500
         
         return jsonify({
             "success": True,
-            "resultUrl": f"/static/{os.path.basename(result_path)}"
+            "resultPath": f"/static/{os.path.basename(result_image)}"
         })
-    
-    except ValueError as e:
-        return jsonify({"success": False, "message": str(e)}), 400
+        
     except Exception as e:
-        return jsonify({"success": False, "message": "Internal server error"}), 500
+        return jsonify({"success": False, "message": f"Error: {str(e)}"}), 500
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)