app.py

import gradio as gr
import os
import json
import shutil
from pathlib import Path
from PIL import Image
import torch
import clip
import numpy as np
import requests
from io import BytesIO
import tempfile

# Global variables for CLIP model (load once, reuse)
_clip_model = None
_clip_preprocess = None
_device = None

def get_clip_model():
    global _clip_model, _clip_preprocess, _device
    if _clip_model is None:
        _device = "cuda" if torch.cuda.is_available() else "cpu"
        _clip_model, _clip_preprocess = clip.load("ViT-B/32", device=_device)
    return _clip_model, _clip_preprocess, _device

class SmartCLIPClassifierNextCloudShare:
    def __init__(self, share_url, share_password, progress_callback=None):
        
        self.share_url = share_url.rstrip('/')
        self.share_password = share_password
        self.progress_callback = progress_callback
        
        self.session = requests.Session()
        self.session.auth = (self.get_share_token(), share_password)
        
        self.temp_dir = tempfile.mkdtemp()
        
        self.categories = [
            "1_Booth",
            "2_Business_Interaction", 
            "3_Buyer_Delegation",
            "4_Aisle",
            "5_Conference",
            "6_Fairground",
            "7_Products",
            "8_Registration",
            "9_Miscellaneous"
        ]
        
        self.log("Loading CLIP model...")
        self.model, self.preprocess, self.device = get_clip_model()
        self.log(f"✅ CLIP loaded on {self.device}")
        
        self.load_deep_analysis()
        
        self.log("🔍 Scanning NextCloud share...")
        self.all_files = self.list_files("")
        self.log(f"Found {len(self.all_files)} total files")
        
        self.get_image_list()
    
    def log(self, message):
        if self.progress_callback:
            self.progress_callback(message)
        print(message)
    
    def get_share_token(self):
        return self.share_url.split('/s/')[-1]
    
    def get_webdav_url(self, path=""):
        token = self.get_share_token()
        base = self.share_url.rsplit('/s/', 1)[0]
        if path:
            return f"{base}/public.php/webdav/{path}"
        return f"{base}/public.php/webdav/"
    
    def download_file(self, filename):
        url = self.get_webdav_url(filename)
        max_retries = 3
        for attempt in range(max_retries):
            try:
                response = self.session.get(url, timeout=60)
                response.raise_for_status()
                return response.content
            except requests.exceptions.Timeout:
                if attempt == max_retries - 1:
                    raise
                self.log(f"Timeout on attempt {attempt + 1}, retrying...")
                continue
    
    def upload_file(self, local_path, remote_filename):
        url = self.get_webdav_url(remote_filename)
        
        # Try to delete existing file first
        try:
            self.session.delete(url, timeout=60)
        except:
            pass
        
        # Now upload with retry
        max_retries = 3
        for attempt in range(max_retries):
            try:
                with open(local_path, 'rb') as f:
                    response = self.session.put(url, data=f, timeout=60)
                response.raise_for_status()
                return
            except requests.exceptions.Timeout:
                if attempt == max_retries - 1:
                    raise
                self.log(f"Upload timeout on attempt {attempt + 1}, retrying...")
                continue
    
    def delete_file(self, filename):
        """Delete a file from NextCloud"""
        url = self.get_webdav_url(filename)
        max_retries = 3
        for attempt in range(max_retries):
            try:
                response = self.session.delete(url, timeout=60)
                response.raise_for_status()
                return True
            except requests.exceptions.Timeout:
                if attempt == max_retries - 1:
                    raise
                self.log(f"Delete timeout on attempt {attempt + 1}, retrying...")
                continue
            except Exception as e:
                self.log(f"Warning: Could not delete {filename}: {e}")
                return False
    
    def list_files(self, remote_path=""):
        url = self.get_webdav_url(remote_path)
        max_retries = 3
        for attempt in range(max_retries):
            try:
                response = self.session.request('PROPFIND', url, headers={'Depth': '1'}, timeout=60)
                response.raise_for_status()
                
                files = []
                lines = response.text.split('<d:href>')
                for line in lines:
                    if '</d:href>' in line:
                        href = line.split('</d:href>')[0]
                        if '/webdav/' in href:
                            filename = href.split('/webdav/')[-1]
                            if filename and not filename.endswith('/'):
                                files.append(filename)
                
                return files
            except requests.exceptions.Timeout:
                if attempt == max_retries - 1:
                    raise
                self.log(f"List files timeout on attempt {attempt + 1}, retrying...")
                continue
    
    def create_folder(self, folder_name):
        url = self.get_webdav_url(folder_name)
        try:
            self.session.request('MKCOL', url, timeout=60)
        except:
            pass
    
    def get_image_list(self):
        self.log("Filtering images...")
        valid_extensions = {'.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp'}
        
        self.images = [f for f in self.all_files if Path(f).suffix.lower() in valid_extensions]
        self.images.sort()
        self.log(f"✅ Found {len(self.images)} images to classify")
    
    def load_deep_analysis(self):
        self.log("Looking for deep_training_analysis.json...")
        
        # Check local file first
        local_json_path = "deep_training_analysis.json"
        
        if os.path.exists(local_json_path):
            self.log(f"Found local file: {local_json_path}")
            try:
                with open(local_json_path, 'r') as f:
                    self.deep_analysis = json.load(f)
                
                self.log("📚 Loaded deep training analysis from local file")
                
                self.category_embeddings = {}
                
                for category in self.categories:
                    if category in self.deep_analysis:
                        data = self.deep_analysis[category]
                        avg_embedding = torch.tensor(data['avg_embedding'], dtype=torch.float32).to(self.device)
                        avg_embedding = avg_embedding / avg_embedding.norm()
                        self.category_embeddings[category] = avg_embedding
                        
                        self.log(f"  {category}: {data['num_training_images']} training images")
                    else:
                        self.category_embeddings[category] = self.create_text_embedding(category)
                
                return
                
            except Exception as e:
                self.log(f"❌ Error loading local deep analysis: {e}")
        
        # Fallback to text embeddings
        self.log("⚠️ deep_training_analysis.json not found - using fallback embeddings")
        self.category_embeddings = {cat: self.create_text_embedding(cat) for cat in self.categories}
    
    def create_text_embedding(self, category):
        descriptions = {
            "1_Booth": "a photo of an exhibition booth at a trade show",
            "2_Business_Interaction": "a photo of business people talking at a trade show",
            "3_Buyer_Delegation": "a photo of a large group visiting a trade show",
            "4_Aisle": "a photo of a trade show aisle between booths",
            "5_Conference": "a photo of a conference presentation or seminar",
            "6_Fairground": "a photo of an exhibition hall or fairground",
            "7_Products": "a photo of products on display",
            "8_Registration": "a photo of a registration desk or entry gate",
            "9_Miscellaneous": "a miscellaneous trade show photo"
        }
        
        text = descriptions.get(category, "a photo")
        text_input = clip.tokenize([text]).to(self.device)
        
        with torch.no_grad():
            text_features = self.model.encode_text(text_input)
            text_features = text_features / text_features.norm(dim=-1, keepdim=True)
        
        return text_features[0]
    
    def classify_image(self, filename):
        try:
            # Download the image
            img_data = self.download_file(filename)
            
            # Classify it
            img = Image.open(BytesIO(img_data)).convert('RGB')
            img_input = self.preprocess(img).unsqueeze(0).to(self.device)
            
            with torch.no_grad():
                img_features = self.model.encode_image(img_input)
                img_features = img_features / img_features.norm(dim=-1, keepdim=True)
                img_features = img_features[0]
                
                similarities = {}
                for category, cat_embedding in self.category_embeddings.items():
                    similarity = (img_features @ cat_embedding).item()
                    similarities[category] = similarity
                
                best_category = max(similarities, key=similarities.get)
                confidence = similarities[best_category]
            
            # Save to temp file
            local_path = os.path.join(self.temp_dir, Path(filename).name)
            with open(local_path, 'wb') as f:
                f.write(img_data)
            
            # Create category folder structure
            category_folder = f"Classified/{best_category}"
            self.create_folder("Classified")
            self.create_folder(category_folder)
            
            # Upload to new location
            remote_dest = f"{category_folder}/{Path(filename).name}"
            self.upload_file(local_path, remote_dest)
            
            # **MOVE instead of COPY: Delete the original file**
            self.delete_file(filename)
            
            # Clean up temp file
            os.remove(local_path)
            
            return best_category, confidence
            
        except Exception as e:
            self.log(f"❌ Error processing {filename}: {str(e)}")
            return "9_Miscellaneous", 0.0
    
    def run(self):
        self.log("🚀 Starting classification...")
        
        # Create category folders
        self.create_folder("Classified")
        for cat in self.categories:
            self.create_folder(f"Classified/{cat}")
        
        stats = {cat: 0 for cat in self.categories}
        confidences = {cat: [] for cat in self.categories}
        
        for i, filename in enumerate(self.images, 1):
            self.log(f"[{i}/{len(self.images)}] Processing {Path(filename).name}...")
            
            category, confidence = self.classify_image(filename)
            
            stats[category] += 1
            confidences[category].append(confidence)
            
            self.log(f"  ✓ Moved to {category} (confidence: {confidence:.3f})")
        
        self.log("✅ CLASSIFICATION COMPLETE!")
        self.log("📁 Results moved to: Classified/")
        self.log("📁 Original files have been removed from root folder")
        
        result_text = "## ✅ Classification Complete!\n\n**Files have been MOVED (not copied) to categorized folders**\n\n**Results Summary:**\n\n"
        for cat in self.categories:
            count = stats[cat]
            if count > 0:
                avg_conf = sum(confidences[cat]) / len(confidences[cat])
                result_text += f"- **{cat}**: {count} images (avg confidence: {avg_conf:.3f})\n"
        
        # Clean up temp directory
        shutil.rmtree(self.temp_dir)
        
        return result_text

def classify_photos(share_url, share_password, progress=gr.Progress()):
    if not share_url or not share_password:
        return "❌ Please enter both the share URL and password", ""
    
    logs = []
    
    def log_callback(message):
        logs.append(message)
        log_display = "\n".join(logs[-20:])
        return log_display
    
    try:
        progress(0, desc="Initializing...")
        
        progress(0.1, desc="Connecting to NextCloud...")
        
        classifier = SmartCLIPClassifierNextCloudShare(
            share_url, 
            share_password,
            progress_callback=log_callback
        )
        
        progress(0.3, desc=f"Found {len(classifier.images)} images to classify")
        
        # Run classification
        result = classifier.run()
        
        progress(1.0, desc="Complete!")
        
        return result, "\n".join(logs)
        
    except requests.exceptions.Timeout:
        error_msg = "❌ Connection Timeout: NextCloud is taking too long to respond.\n\nPlease check your network connection and try again."
        return error_msg, "\n".join(logs)
    except requests.exceptions.RequestException as e:
        error_msg = f"❌ Connection Error: Could not connect to NextCloud.\n\nPlease check:\n- Your share URL is correct\n- Your password is correct\n- The share link has 'Allow upload and editing' enabled\n\nError details: {str(e)}"
        return error_msg, "\n".join(logs)
    except Exception as e:
        error_msg = f"❌ Error: {str(e)}\n\nPlease check your share URL and password and try again."
        return error_msg, "\n".join(logs)

# Gradio Interface
with gr.Blocks(title="Trade Show Photo Classifier", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🤖 Trade Show Photo Classifier")
    gr.Markdown("Automatically classify your trade show photos using AI-powered image recognition")
    
    gr.Markdown("""
    ### 📋 Setup Instructions:
    
    1. Upload your photos to a NextCloud folder
    2. Create a public share link for that folder
    3. **Important:** When creating the share, enable **"Allow upload and editing"** (click the three dots → Share settings)
    4. Set a password for the share
    5. Copy the share URL and password below
    6. Click "Start Classification"
    
    **⚠️ Note: Files will be MOVED (not copied) to the Classified folder. Original files will be deleted from the root.**
    """)
    
    with gr.Row():
        with gr.Column():
            share_url = gr.Textbox(
                label="NextCloud Share URL",
                placeholder="https://cloud2.messefrankfurtexchange.com/s/...",
                info="Enter the public share link to your NextCloud folder"
            )
            share_password = gr.Textbox(
                label="Share Password",
                type="password",
                info="Enter the password for the NextCloud share"
            )
            classify_btn = gr.Button("🚀 Start Classification", variant="primary", size="lg")
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### 📊 Results")
            output = gr.Markdown()
        
        with gr.Column():
            gr.Markdown("### 📝 Classification Log")
            logs_output = gr.Textbox(
                label="Progress",
                lines=15,
                max_lines=20,
                interactive=False,
                show_label=False
            )
    
    classify_btn.click(
        fn=classify_photos,
        inputs=[share_url, share_password],
        outputs=[output, logs_output]
    )
    
    gr.Markdown("---")
    gr.Markdown("""
    **Categories:**
    - 1_Booth: Exhibition booths
    - 2_Business_Interaction: Business conversations
    - 3_Buyer_Delegation: Group visits
    - 4_Aisle: Walkways between booths
    - 5_Conference: Presentations and seminars
    - 6_Fairground: Exhibition halls
    - 7_Products: Product displays
    - 8_Registration: Entry and registration areas
    - 9_Miscellaneous: Other trade show content
    
    *Powered by OpenAI CLIP | Deployed on Hugging Face Spaces*
    """)

if __name__ == "__main__":
    demo.launch()