modifications and added sam_vit_b.pth

app.py

@@ -2,136 +2,108 @@ import streamlit as st
 import numpy as np
 import cv2
 from PIL import Image
-import torch
 import torchvision.transforms as transforms
-import torchvision.models as models
+from segment_anything import sam_model_registry, SamPredictor
 from io import BytesIO
-from google.generativeai import configure, GenerativeModel
 import base64
+from google.generativeai import configure, GenerativeModel
 
 # Configure Gemini API
 configure(api_key="AIzaSyBawh403z5cyyQzFhQo14y7oUQw6nr8mIg")
 model = GenerativeModel("gemini-2.0-flash")
 
-# Load DeepLabV3 model for garment segmentation
-def load_segmentation_model():
-    model = models.segmentation.deeplabv3_resnet101(pretrained=True)
-    model.eval()
-    return model
-
-def segment_garment(image, model):
-    # Convert to model-compatible format
-    transform = transforms.Compose([
-        transforms.Resize((520, 520)),  
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-    ])
-    
-    image_tensor = transform(image).unsqueeze(0)
-
-    with torch.no_grad():
-        output = model(image_tensor)['out'][0]
-
-    # Convert output to segmentation mask
-    mask = output.argmax(0).byte().cpu().numpy()
+# Load SAM model with ViT-Base
+def load_sam_model():
+    sam = sam_model_registry["vit_b"](checkpoint="sam_vit_b.pth")  # Use vit_b instead of vit_h
+    predictor = SamPredictor(sam)
+    return predictor
 
-    # Debugging: Check if any garment pixels are detected
-    print("Unique values in mask:", np.unique(mask))
+# Preprocess image using OpenCV (Edge Detection & Background Removal)
+def preprocess_image(image):
+    image = np.array(image.convert("RGB"))  # Convert to NumPy array
+    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)  # Convert to Grayscale
+    edges = cv2.Canny(gray, 100, 200)  # Apply Canny Edge Detection
+    return Image.fromarray(edges)  # Convert back to PIL Image
 
-    if mask.max() == 0:  # No garment detected
-        print("Warning: No garment detected. Check input image.")
-        return image  # Return original image
+# Segment garment using SAM
+def segment_garment(image, predictor):
+    image_np = np.array(image.convert("RGB"))
+    predictor.set_image(image_np)
 
-    # Resize mask to match original image size
-    mask = Image.fromarray(mask.astype(np.uint8) * 255)
-    mask = mask.resize(image.size, Image.NEAREST)  # Match original size
-    mask = np.array(mask)
+    # Use center point of image as prompt
+    height, width, _ = image_np.shape
+    input_point = np.array([[width // 2, height // 2]])
+    input_label = np.array([1])  # 1 indicates object selection
 
-    # Convert grayscale mask to 3-channel
-    mask = np.stack([mask] * 3, axis=-1)  
+    masks, _, _ = predictor.predict(point_coords=input_point, point_labels=input_label)
+    mask = masks[0]  # Get first mask
 
-    # Convert original image to numpy array
-    image_np = np.array(image)
+    # Resize mask to match image
+    mask_resized = cv2.resize(mask.astype(np.uint8) * 255, (width, height), interpolation=cv2.INTER_NEAREST)
+    mask_resized = np.stack([mask_resized] * 3, axis=-1)  # Convert to 3-channel
 
-    # Apply the mask (only keep garment pixels)
-    segmented = np.where(mask > 0, image_np, 0)
+    # Apply segmentation mask
+    segmented = np.where(mask_resized > 0, image_np, 0)
 
     return Image.fromarray(segmented)
 
-# Preprocess image using OpenCV (Edge Detection & Background Removal)
-def preprocess_image(image):
-    image = np.array(image.convert("RGB"))
-    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
-    edges = cv2.Canny(gray, 100, 200)
-    return Image.fromarray(edges)
-
-# AI garment analysis & fashion recommendations
+# AI garment analysis
 def analyze_garment(image):
-    # Convert PIL Image to bytes
     image_bytes = BytesIO()
     image.save(image_bytes, format="PNG")
-    image_bytes = image_bytes.getvalue()
-
-    # Encode image in Base64
-    encoded_image = base64.b64encode(image_bytes).decode("utf-8")
+    encoded_image = base64.b64encode(image_bytes.getvalue()).decode("utf-8")
 
-    # Prepare request payload in correct Gemini API format
     prompt = {
         "parts": [
             {"text": "Analyze the garment in this image, describing its style, fabric, and design elements. "
-                     "Based on the garment's features, suggest the best occasions to wear it "
-                     "and recommend complementary fashion pieces (e.g., shoes, accessories, layering options). "
-                     "Also, provide a seasonal suitability rating."},
+                     "Suggest the best occasions to wear it and recommend complementary fashion pieces."},
             {"inline_data": {"mime_type": "image/png", "data": encoded_image}}
         ]
     }
-
-    # Call Gemini API
+    
     response = model.generate_content(prompt)
-
     return response.text if response else "Analysis failed."
 
-# Load segmentation model
-segmentation_model = load_segmentation_model()
+# Load SAM model
+sam_predictor = load_sam_model()
 
 # Streamlit UI
-st.title("👗 AI-Enhanced Fashion Design with Gemini 2.0 Flash")
-st.write("""
-### 🎨 **About This App**
-This AI-powered fashion tool helps you analyze and enhance garment designs using Google's **Gemini 2.0 Flash** and DeepLabV3 segmentation.
-It provides recommendations on **suitable fashion styles** for different occasions based on the uploaded garment.
-
-### 🛠 **How to Use This App**
-1. **Upload an Image** – Select a clothing item or a fashion sketch.
-2. **View Preprocessing & Segmentation** – The app applies edge detection and garment segmentation.
-3. **Analyze Garment** – Click "Analyze Garment" to get AI-powered fashion insights.
-4. **Get Recommendations** – The AI suggests suitable **occasions and styling tips** based on the garment.
-
-### 🔍 **How It Works**
-- **Edge Detection**: Uses OpenCV to highlight contours and details in the garment.
-- **Garment Segmentation**: DeepLabV3 identifies the clothing item and removes the background.
-- **AI Fashion Analysis**: Google’s Gemini AI analyzes the **style, fabric, and design** of the garment and provides recommendations.
-
-➡️ Try it now by uploading an image of your clothing!
+st.title("👗 AI Fashion Analysis with SAM & Gemini AI")
+
+# Description and Instructions
+st.markdown("""
+### 📌 How to Use this App:
+1. *Upload an Image*: Click the upload button and select a fashion image.
+2. *Preprocess the Image*: Click 'Preprocess' to apply edge detection and garment segmentation.
+3. *View Results*: Processed and segmented images will be displayed.
+4. *Analyze the Garment*: Click 'Analyze Garment' to get AI-based fashion insights.
 """)
 
 # File Upload
 uploaded_file = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])
 
 if uploaded_file is not None:
-    # Display uploaded image
     image = Image.open(uploaded_file)
     st.image(image, caption="Uploaded Image", use_container_width=True)
-    
-    # Preprocess Image with OpenCV
-    processed_image = preprocess_image(image)
-    st.image(processed_image, caption="Processed Image (Edge Detection)", use_container_width=True)
-    
-    # Segment Garment using DeepLabV3
-    segmented_image = segment_garment(image, segmentation_model)
-    st.image(segmented_image.convert("RGB"), caption="Segmented Garment", use_container_width=True)
 
-    
-if st.button("Analyze Garment"):
-    result = analyze_garment(image)
-    st.success(result)
+    # Initialize session state for persistence
+    if "processed_image" not in st.session_state:
+        st.session_state.processed_image = None
+    if "segmented_image" not in st.session_state:
+        st.session_state.segmented_image = None
+
+    # Preprocess Button (Runs Edge Detection & Segmentation)
+    if st.button("Preprocess"):
+        st.session_state.processed_image = preprocess_image(image)
+        st.session_state.segmented_image = segment_garment(image, sam_predictor)
+
+    # Display Preprocessed Images if Available
+    if st.session_state.processed_image:
+        st.image(st.session_state.processed_image, caption="Edge Detection", use_container_width=True)
+    if st.session_state.segmented_image:
+        st.image(st.session_state.segmented_image, caption="Segmented Garment", use_container_width=True)
+
+    # Analyze Garment Button (Gemini AI)
+    if st.button("Analyze Garment"):
+        result = analyze_garment(image)
+        st.success(result)

requirements.txt

@@ -4,4 +4,5 @@ opencv-python
 pillow
 torch
 torchvision
-google-generativeai
+google-generativeai
+transformers

sam_vit_b.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ec2df62732614e57411cdcf32a23ffdf28910380d03139ee0f4fcbe91eb8c912
+size 375042383