Update app.py

app.py

@@ -2,11 +2,6 @@ import os
 import cv2
 import gradio as gr
 import numpy as np
-import random
-import base64
-import requests
-import json
-import time
 from transformers import DetrForObjectDetection, DetrImageProcessor
 import torch
 
@@ -25,7 +20,7 @@ def detect_face_and_neck(image):
     face_box = None
     for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
         if score > 0.7:
-            if label == 1:  # Person
+            if label == 1:  # Person (this can include neck)
                 neck_box = box
             elif label == 2:  # Face
                 face_box = box
@@ -33,24 +28,37 @@ def detect_face_and_neck(image):
     return face_box, neck_box
 
 # Function to overlay jewelry on the detected regions
-def place_jewelry(image, jewelry_image, jewelry_type, position):
+def place_jewelry(image, jewelry_image, position):
     x, y, w, h = position
-    resized_jewelry = cv2.resize(jewelry_image, (w, h))
+    resized_jewelry = cv2.resize(jewelry_image, (int(w), int(h)))
     
-    image[y:y+h, x:x+w] = resized_jewelry
+    # Ensure that the image has an alpha channel (RGBA) for blending
+    if resized_jewelry.shape[2] == 4:
+        # Blending using alpha transparency
+        for c in range(0, 3):
+            image[y:y+h, x:x+w, c] = resized_jewelry[:, :, c] * (resized_jewelry[:, :, 3] / 255.0) + image[y:y+h, x:x+w, c] * (1.0 - resized_jewelry[:, :, 3] / 255.0)
+    else:
+        image[y:y+h, x:x+w] = resized_jewelry
     
     return image
 
 # Try-on function for jewelry
 def tryon_jewelry(person_img, jewelry_img, jewelry_type):
+    # Ensure images are valid
+    if person_img is None or jewelry_img is None:
+        return None
+    
+    # Detect face and neck using Hugging Face model
     face_box, neck_box = detect_face_and_neck(person_img)
-
+    
     if jewelry_type == "Necklace" and neck_box is not None:
-        result_img = place_jewelry(person_img, jewelry_img, "Necklace", neck_box)
+        # Apply necklace on neck region
+        result_img = place_jewelry(person_img, jewelry_img, neck_box)
     elif jewelry_type == "Earrings" and face_box is not None:
-        result_img = place_jewelry(person_img, jewelry_img, "Earrings", face_box)
+        # Assuming ears are part of the face box for simplicity
+        result_img = place_jewelry(person_img, jewelry_img, face_box)
     else:
-        result_img = person_img  # Return original image if no detection
+        result_img = person_img  # If no detection, return original image
 
     return result_img