Update app.py

app.py

@@ -1,119 +1,199 @@
-# IMPORTANT: These environment variables must be set before any other imports
 import os
 os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
 os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION_VERSION"] = "2"
 
 import streamlit as st
-import cv2
-import numpy as np
+import torch
 from PIL import Image
-from mtcnn import MTCNN
-from deepface import DeepFace
-from transformers import BlipProcessor, BlipForConditionalGeneration
-
-# Initialize MTCNN for face detection
-detector = MTCNN()
-
-# Load BLIP model for scene description
+import numpy as np
+from transformers import (
+    BlipProcessor, 
+    BlipForConditionalGeneration,
+    AutoProcessor, 
+    AutoModelForObjectDetection,
+    CLIPProcessor, 
+    CLIPModel,
+    ViTImageProcessor, 
+    ViTForImageClassification
+)
+
+# Initialize models with caching
 @st.cache_resource
-def load_blip_model():
-    processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-    model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
-    return processor, model
-
-processor, blip_model = load_blip_model()
-
-# Function for face analysis using DeepFace
-def analyze_face(face_image):
-    face_array = np.array(face_image)
-    try:
-        result = DeepFace.analyze(face_array, actions=['emotion', 'gender', 'age', 'race'], enforce_detection=False)
-        return result
-    except Exception as e:
-        return {"error": str(e)}
-
-# Function for glasses detection using Haar Cascade
-def detect_glasses(face_image):
-    # Path to Haar Cascade in the main folder
-    haar_cascade_path = os.path.join(os.getcwd(), 'haarcascade_frontalface_default.xml')
+def load_models():
+    """Load and cache all required models"""
+    
+    # BLIP for detailed image captioning
+    blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
+    blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large")
     
-    # Load Haar Cascade for glasses detection
-    glasses_cascade = cv2.CascadeClassifier(haar_cascade_path)
+    # DETR for object detection
+    detr_processor = AutoProcessor.from_pretrained("facebook/detr-resnet-50")
+    detr_model = AutoModelForObjectDetection.from_pretrained("facebook/detr-resnet-50")
     
-    # Convert to grayscale
-    face_array = np.array(face_image.convert("L"))
+    # CLIP for visual attributes and concepts
+    clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+    clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
     
-    # Perform detection
-    detected_glasses = glasses_cascade.detectMultiScale(face_array, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))
+    # ViT for scene classification
+    vit_processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224')
+    vit_model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
     
-    return "Yes" if len(detected_glasses) > 0 else "No"
+    return (
+        blip_processor, blip_model,
+        detr_processor, detr_model,
+        clip_processor, clip_model,
+        vit_processor, vit_model
+    )
 
-# Function for scene description using BLIP
 @st.cache_data
-def describe_scene(image):
-    inputs = processor(image, return_tensors="pt")
-    outputs = blip_model.generate(**inputs)
+def generate_caption(_image, processor, model):
+    """Generate detailed image caption using BLIP"""
+    inputs = processor(_image, return_tensors="pt")
+    outputs = model.generate(
+        **inputs,
+        max_length=100,
+        num_beams=5,
+        temperature=1.0,
+        top_p=0.9,
+        repetition_penalty=1.5
+    )
     caption = processor.decode(outputs[0], skip_special_tokens=True)
     return caption
 
-# Main Streamlit application
-def main():
-    st.title("Comprehensive Image Analysis")
-    st.write("Upload an image to analyze faces, detect glasses, and provide scene descriptions.")
-
-    uploaded_image = st.file_uploader("Upload an image...", type=["jpg", "jpeg", "png"])
-
-    if uploaded_image is not None:
-        image = Image.open(uploaded_image).convert("RGB")
-        st.image(image, caption="Uploaded Image", use_container_width=True)
-
-        if st.button("Analyze"):
-            with st.spinner("Analyzing..."):
-                # Scene description using BLIP
-                scene_caption = describe_scene(image)
-                st.subheader("Scene Description")
-                st.write(scene_caption)
-
-                # Detect faces using MTCNN
-                faces = detector.detect_faces(np.array(image))
-
-                if faces:
-                    st.write(f"**Number of Faces Detected:** {len(faces)}")
-
-                    for i, face in enumerate(faces):
-                        x, y, width, height = face['box']
-                        face_crop = image.crop((x, y, x + width, y + height))
-                        face_crop_resized = face_crop.resize((150, 150))
-
-                        st.subheader(f"Analysis for Face {i + 1}")
-                        st.image(face_crop_resized, caption=f"Face {i + 1}", width=150)
-
-                        # Analyze face with DeepFace
-                        result = analyze_face(face_crop)
-                        if "error" in result:
-                            st.error(f"Face Analysis Error: {result['error']}")
-                        else:
-                            single_result = result[0] if isinstance(result, list) else result
-                            gender = single_result.get('gender', 'N/A')
-                            age = single_result.get('age', 'N/A')
-                            race = single_result.get('dominant_race', 'N/A')
-                            emotions = single_result.get('emotion', {})
+@st.cache_data
+def detect_objects(_image, processor, model):
+    """Detect objects in the image using DETR"""
+    inputs = processor(_image, return_tensors="pt")
+    outputs = model(**inputs)
+    
+    # Post-process object detection results
+    target_sizes = torch.tensor([_image.size[::-1]])
+    results = processor.post_process_object_detection(
+        outputs, 
+        target_sizes=target_sizes,
+        threshold=0.7
+    )[0]
+    
+    objects = []
+    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+        objects.append({
+            'label': model.config.id2label[label.item()],
+            'confidence': score.item(),
+            'box': box.tolist()
+        })
+    
+    return objects
 
-                            st.write(f"**Gender:** {gender}")
-                            st.write(f"**Age:** {age}")
-                            st.write(f"**Race:** {race}")
+@st.cache_data
+def analyze_visual_attributes(_image, processor, model):
+    """Analyze visual attributes using CLIP"""
+    # Define candidate attributes to check
+    attributes = [
+        "bright", "dark", "colorful", "monochrome",
+        "natural", "artificial", "indoor", "outdoor",
+        "modern", "vintage", "abstract", "realistic",
+        "simple", "complex", "symmetrical", "asymmetrical",
+        "high contrast", "low contrast", "textured", "smooth"
+    ]
+    
+    inputs = processor(
+        images=_image,
+        text=attributes,
+        return_tensors="pt",
+        padding=True
+    )
+    
+    outputs = model(**inputs)
+    probs = outputs.logits_per_image.softmax(dim=1)[0]
+    
+    return [(attr, prob.item()) for attr, prob in zip(attributes, probs)]
 
-                            # Display emotions as text with percentages
-                            if emotions:
-                                sorted_emotions = sorted(emotions.items(), key=lambda x: x[1], reverse=True)
-                                emotion_text = ", ".join([f"{emotion}: {prob:.2f}%" for emotion, prob in sorted_emotions])
-                                st.write(f"**Emotions:** {emotion_text}")
+@st.cache_data
+def classify_scene(_image, processor, model):
+    """Classify the scene type using ViT"""
+    inputs = processor(_image, return_tensors="pt")
+    outputs = model(**inputs)
+    probs = outputs.logits.softmax(dim=-1)[0]
+    
+    # Get top 5 predictions
+    top_probs, top_indices = torch.topk(probs, 5)
+    predictions = [
+        (model.config.id2label[idx.item()], prob.item())
+        for idx, prob in zip(top_indices, top_probs)
+    ]
+    
+    return predictions
+
+def format_attributes(attributes):
+    """Format visual attributes for display"""
+    significant_attrs = [(attr, score) for attr, score in attributes if score > 0.7]
+    if significant_attrs:
+        return ", ".join([
+            f"{attr} ({score:.1%})" 
+            for attr, score in sorted(significant_attrs, key=lambda x: x[1], reverse=True)
+        ])
+    return "No strong visual attributes detected"
 
-                            # Glasses detection
-                            glasses = detect_glasses(face_crop)
-                            st.write(f"**Glasses:** {glasses}")
-                else:
-                    st.write("No faces detected.")
+def main():
+    st.set_page_config(page_title="Image Description System", layout="wide")
+    st.title("🖼️ Comprehensive Image Description System")
+    st.write("Upload an image to get a detailed analysis of its content and characteristics.")
+    
+    # Load all models
+    with st.spinner("Loading models..."):
+        (
+            blip_processor, blip_model,
+            detr_processor, detr_model,
+            clip_processor, clip_model,
+            vit_processor, vit_model
+        ) = load_models()
+    
+    # File uploader
+    uploaded_file = st.file_uploader(
+        "Choose an image...", 
+        type=['jpg', 'jpeg', 'png'],
+        help="Upload a clear image for the best analysis results"
+    )
+    
+    if uploaded_file:
+        # Load and display image
+        image = Image.open(uploaded_file).convert('RGB')
+        st.image(image, caption="Uploaded Image", use_column_width=True)
+        
+        if st.button("🔍 Analyze Image"):
+            with st.spinner("Analyzing image..."):
+                # 1. Generate detailed caption
+                caption = generate_caption(image, blip_processor, blip_model)
+                st.subheader("📝 Detailed Description")
+                st.write(caption)
+                
+                # Create two columns for remaining analysis
+                col1, col2 = st.columns(2)
+                
+                with col1:
+                    # 2. Detect and list objects
+                    st.subheader("🎯 Detected Objects")
+                    objects = detect_objects(image, detr_processor, detr_model)
+                    if objects:
+                        for obj in sorted(objects, key=lambda x: x['confidence'], reverse=True):
+                            st.write(f"• {obj['label']} ({obj['confidence']:.1%} confidence)")
+                    else:
+                        st.write("No objects detected with high confidence.")
+                    
+                    # 3. Scene classification
+                    st.subheader("🏞️ Scene Classification")
+                    scenes = classify_scene(image, vit_processor, vit_model)
+                    for scene, prob in scenes:
+                        st.write(f"• {scene} ({prob:.1%} confidence)")
+                
+                with col2:
+                    # 4. Visual attributes
+                    st.subheader("🎨 Visual Attributes")
+                    attributes = analyze_visual_attributes(image, clip_processor, clip_model)
+                    formatted_attrs = format_attributes(attributes)
+                    st.write(formatted_attrs)
+                
+                st.success("Analysis complete! ✨")
 
 if __name__ == "__main__":
     main()