gaur3009's picture
Update app.py
6246ec7 verified
raw
history blame
1.74 kB
import gradio as gr
import torch
from torchvision import models, transforms
from PIL import Image, ImageEnhance
import numpy as np
import cv2
model = models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
model.eval()
transform = transforms.Compose([
transforms.ToTensor()
])
def detect_dress(image):
image_tensor = transform(image).unsqueeze(0)
with torch.no_grad():
outputs = model(image_tensor)
boxes = outputs[0]['boxes'].numpy()
scores = outputs[0]['scores'].numpy()
threshold = 0.8
dress_boxes = [box for box, score in zip(boxes, scores) if score > threshold]
draw = ImageDraw.Draw(image)
for box in dress_boxes:
draw.rectangle(box.tolist(), outline="red", width=3)
return image, dress_boxes
def crop_image(image, box):
return image.crop(box)
def adjust_color(image, factor):
enhancer = ImageEnhance.Color(image)
return enhancer.enhance(factor)
def process_image(image, edit_type, factor):
detected_image, boxes = detect_dress(image)
if not boxes:
return detected_image, "No dresses detected."
if edit_type == "Crop":
box = boxes[0]
edited_image = crop_image(image, box)
elif edit_type == "Adjust Color":
edited_image = adjust_color(image, factor)
else:
edited_image = image
return edited_image, "Edit applied."
iface = gr.Interface(
fn=process_image,
inputs=[
gr.inputs.Image(type="pil"),
gr.inputs.Radio(choices=["None", "Crop", "Adjust Color"]),
gr.inputs.Slider(0.5, 2.0, step=0.1, label="Factor")
],
outputs=[
gr.outputs.Image(type="pil"),
gr.outputs.Textbox(label="Result")
],
live=True
)
iface.launch()