import torch import torch.nn as nn from torchvision import models import torchvision.models.detection as detection import torchvision.transforms as transforms from PIL import Image import segmentation_models_pytorch as smp import numpy as np import cv2 import gradio as gr from pytorch_grad_cam import GradCAM from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget from pytorch_grad_cam.utils.image import show_cam_on_image import os import logging # Set up logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) # Device setup (use CPU for Spaces free tier) device = torch.device('cpu') # Transforms for segmentation (256x256) seg_transform = transforms.Compose([ transforms.Resize((256, 256)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) # Transforms for classification (224x224) preprocess = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) # Pipeline function (load models lazily) def process_image(input_image): try: logger.info("Loading U-Net model...") seg_model = smp.Unet( encoder_name='resnet34', classes=1, activation=None ) seg_model.load_state_dict(torch.load('best_dress_segmentation_model.pth', map_location=device)) seg_model.to(device) seg_model.eval() logger.info("U-Net loaded") logger.info("Loading ResNet18 model...") class_list = ['casual_dress', 'denim_dress', 'evening_dress', 'jersey_dress', 'knitted_dress', 'maxi_dress', 'occasion_dress', 'shift_dress', 'shirt_dress', 'work_dress'] cnn_model = models.resnet18(weights=None) num_ftrs = cnn_model.fc.in_features cnn_model.fc = nn.Linear(num_ftrs, len(class_list)) cnn_model.load_state_dict(torch.load('final_model.pth', map_location=device)) cnn_model.to(device) cnn_model.eval() logger.info("ResNet18 loaded") logger.info("Loading Faster R-CNN model...") human_detector = detection.fasterrcnn_resnet50_fpn(weights='DEFAULT') human_detector.eval().to(device) logger.info("Faster R-CNN loaded") logger.info("Processing image...") img = Image.fromarray(input_image).convert('RGB') width, height = img.size input_tensor = seg_transform(img).unsqueeze(0).to(device) with torch.no_grad(): predictions = human_detector([input_tensor]) pred = predictions[0] boxes = pred['boxes'] labels = pred['labels'] scores = pred['scores'] # Filter for humans (label 1) with score > 0.2 human_boxes = [] human_scores = [] for box, label, score in zip(boxes, labels, scores): if label == 1 and score > 0.2: # Lowered threshold x1, y1, x2, y2 = box.tolist() human_boxes.append([x1, y1, x2, y2]) human_scores.append(score) if not human_boxes: logger.warning("No humans detected") return "No humans detected.", "No dress segmented.", "No category predicted.", "No Grad-CAM available." # Use the most confident box max_score_idx = torch.argmax(torch.tensor(human_scores)).item() bbox = human_boxes[max_score_idx] logger.info(f"Selected bounding box: {bbox}") # Crop crop_img = img.crop((int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3]))) # Segmentation logger.info("Running segmentation...") input_tensor = seg_transform(crop_img).unsqueeze(0).to(device) with torch.no_grad(): seg_output = seg_model(input_tensor).sigmoid().cpu().squeeze() seg_mask = (seg_output > 0.5).float() if seg_mask.sum() == 0: logger.warning("No dress detected") return "No dress detected.", "No dress segmented.", "No category predicted.", "No Grad-CAM available." seg_mask_pil = transforms.ToPILImage()(seg_mask.unsqueeze(0)) seg_mask_pil = seg_mask_pil.resize(crop_img.size, Image.Resampling.BILINEAR) seg_mask_np = np.array(seg_mask_pil) / 255.0 crop_img_np = np.array(crop_img) masked_img_np = crop_img_np * np.expand_dims(seg_mask_np, axis=2) masked_img = Image.fromarray(masked_img_np.astype(np.uint8)) # Classification logger.info("Running classification...") input_tensor = preprocess(masked_img).unsqueeze(0).to(device) with torch.no_grad(): output = cnn_model(input_tensor) pred_class = torch.argmax(output, dim=1).item() category = class_list[pred_class] logger.info(f"Predicted category: {category}") # Grad-CAM logger.info("Generating Grad-CAM...") target_layers = [cnn_model.layer4[-1]] cam = GradCAM(model=cnn_model, target_layers=target_layers) targets = [ClassifierOutputTarget(pred_class)] grayscale_cam = cam(input_tensor=input_tensor, targets=targets)[0, :] visualization = show_cam_on_image(masked_img_np / 255.0, grayscale_cam, use_rgb=True, colormap=cv2.COLORMAP_JET) # Draw bounding box on original image original_with_box = np.array(img) cv2.rectangle(original_with_box, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), (255, 0, 0), 2) logger.info("Processing complete") return Image.fromarray(original_with_box), masked_img, category, Image.fromarray(visualization) except Exception as e: logger.error(f"Error processing image: {str(e)}") return f"Error processing image: {str(e)}", "Error.", "Error.", "Error." # Gradio interface logger.info("Initializing Gradio interface...") with gr.Blocks() as demo: gr.Markdown("### Dress Classification Pipeline\nUpload an image to detect humans, segment dresses, classify, and visualize Grad-CAM.") image_input = gr.Image(label="Upload Image") original_output = gr.Image(label="Original Image with Bounding Box") segmented_output = gr.Image(label="Segmented Dress") category_output = gr.Textbox(label="Predicted Dress Category") gradcam_output = gr.Image(label="Grad-CAM Heatmap") image_input.change(process_image, image_input, [original_output, segmented_output, category_output, gradcam_output]) logger.info("Launching Gradio app...") demo.launch(server_name="0.0.0.0", server_port=7860)