Update app.py

app.py

@@ -4,22 +4,28 @@ from torchvision import models, transforms
 from PIL import Image, ImageEnhance, ImageDraw
 import numpy as np
 
+# Load pre-trained Faster R-CNN model with updated weights parameter
 model = models.detection.fasterrcnn_resnet50_fpn(weights=models.detection.FasterRCNN_ResNet50_FPN_Weights.COCO_V1)
 model.eval()
 
+# Image transformation
 transform = transforms.Compose([
     transforms.ToTensor()
 ])
 
+# Detection function
 def detect_dress(image):
     image_tensor = transform(image).unsqueeze(0)
     with torch.no_grad():
         outputs = model(image_tensor)
     
     boxes = outputs[0]['boxes'].numpy()
+    labels = outputs[0]['labels'].numpy()
     scores = outputs[0]['scores'].numpy()
-    threshold = 0.8
-    dress_boxes = [box for box, score in zip(boxes, scores) if score > threshold]
+    
+    # Filter for dress class (e.g., class 1 in COCO for 'person' which might also work for clothing)
+    threshold = 0.5
+    dress_boxes = [box for box, label, score in zip(boxes, labels, scores) if label == 1 and score > threshold]
 
     draw = ImageDraw.Draw(image)
     for box in dress_boxes:
@@ -27,6 +33,7 @@ def detect_dress(image):
     
     return image, dress_boxes
 
+# Image editing functions
 def crop_image(image, box):
     return image.crop(box)
 
@@ -34,6 +41,7 @@ def adjust_color(image, factor):
     enhancer = ImageEnhance.Color(image)
     return enhancer.enhance(factor)
 
+# Gradio interface function
 def process_image(image, edit_type, factor):
     detected_image, boxes = detect_dress(image)
     
@@ -50,6 +58,7 @@ def process_image(image, edit_type, factor):
 
     return edited_image, "Edit applied."
 
+# Create Gradio interface
 iface = gr.Interface(
     fn=process_image,
     inputs=[
@@ -64,4 +73,5 @@ iface = gr.Interface(
     live=True
 )
 
+# Launch the Gradio interface
 iface.launch()