Update app.py

app.py

@@ -1,159 +1,53 @@
 import torch
 import os
-import glob
 from PIL import Image
-from transformers import AutoProcessor, AutoModelForImageClassification
+from transformers import AutoModelForImageClassification, SiglipImageProcessor
 import gradio as gr
 import pytesseract
 
-def find_model_files():
-    """Find model files in the current directory structure"""
-    print("=== Searching for model files ===")
-    
-    # Look for key model files
-    config_files = glob.glob("**/config.json", recursive=True)
-    model_files = glob.glob("**/pytorch_model.bin", recursive=True) + glob.glob("**/model.safetensors", recursive=True)
-    preprocessor_files = glob.glob("**/preprocessor_config.json", recursive=True)
-    
-    print(f"Found config.json files: {config_files}")
-    print(f"Found model weight files: {model_files}")
-    print(f"Found preprocessor_config.json files: {preprocessor_files}")
-    
-    # Find the directory that contains all necessary files
-    for config_file in config_files:
-        model_dir = os.path.dirname(config_file)
-        if not model_dir:  # If config.json is in root
-            model_dir = "."
-        
-        # Check if this directory has all required files
-        has_model = any(os.path.dirname(f) == model_dir or (not os.path.dirname(f) and model_dir == ".") for f in model_files)
-        has_preprocessor = any(os.path.dirname(f) == model_dir or (not os.path.dirname(f) and model_dir == ".") for f in preprocessor_files)
-        
-        if has_model and has_preprocessor:
-            print(f"Found complete model in directory: {model_dir}")
-            return model_dir
-        elif has_model:
-            print(f"Found model with config but missing preprocessor in: {model_dir}")
-            return model_dir  # Try anyway, might work
-    
-    print("No complete model directory found")
-    return None
-
-# Search for model files
-MODEL_PATH = find_model_files()
-if MODEL_PATH is None:
-    MODEL_PATH = "."  # Fallback to current directory
-    print("Falling back to current directory")
+# Model path
+MODEL_PATH = "./model"
 
 try:
-    # Load model and processor from detected path
-    print(f"=== Attempting to load model from: {MODEL_PATH} ===")
-    print(f"Current working directory: {os.getcwd()}")
+    print(f"=== Loading model from: {MODEL_PATH} ===")
+    print(f"Available files: {os.listdir(MODEL_PATH)}")
     
-    # List all files in the detected model directory
-    if MODEL_PATH == ".":
-        print("Files in root directory:")
-        for item in os.listdir("."):
-            if os.path.isfile(item):
-                print(f"  File: {item}")
-            else:
-                print(f"  Directory: {item}/")
-                try:
-                    sub_files = os.listdir(item)[:5]  # Show first 5 files
-                    print(f"    Contains: {sub_files}{'...' if len(os.listdir(item)) > 5 else ''}")
-                except:
-                    pass
-    else:
-        print(f"Files in {MODEL_PATH}:")
-        print(f"  {os.listdir(MODEL_PATH)}")
-    
-    # Try to load the model
+    # Load the model (this should work with your files)
     print("Loading model...")
     model = AutoModelForImageClassification.from_pretrained(MODEL_PATH, local_files_only=True)
-    print("Model loaded successfully!")
+    print("✅ Model loaded successfully!")
     
-    print("Loading processor...")
+    # Load just the image processor (not the full AutoProcessor)
+    print("Loading image processor...")
     try:
-        processor = AutoProcessor.from_pretrained(MODEL_PATH, local_files_only=True)
-        print("Processor loaded successfully!")
-    except Exception as proc_error:
-        print(f"Error loading processor from local files: {proc_error}")
-        print("Attempting to load just the image processor...")
-        
-        # Try to load just the image processor from your model
-        try:
-            from transformers import SiglipImageProcessor
-            processor = SiglipImageProcessor.from_pretrained(MODEL_PATH, local_files_only=True)
-            print("Image processor loaded successfully from local files!")
-        except Exception as img_proc_error:
-            print(f"Error loading local image processor: {img_proc_error}")
-            print("Loading image processor from base SigLIP model...")
-            
-            # Try to load processor from the base SigLIP model
-            try:
-                from transformers import SiglipImageProcessor
-                processor = SiglipImageProcessor.from_pretrained("google/siglip-base-patch16-224")
-                print("Image processor loaded from base SigLIP model!")
-            except Exception as base_error:
-                print(f"Error loading base processor: {base_error}")
-                print("Using CLIP processor as fallback...")
-                
-                # As a last resort, try to create a minimal processor
-                from transformers import CLIPImageProcessor
-                processor = CLIPImageProcessor.from_pretrained("openai/clip-vit-base-patch32")
-                print("Using CLIP processor as fallback!")
+        # Try to load the image processor from your local files
+        processor = SiglipImageProcessor.from_pretrained(MODEL_PATH, local_files_only=True)
+        print("✅ Image processor loaded from local files!")
+    except Exception as e:
+        print(f"⚠️  Could not load local processor: {e}")
+        print("Loading image processor from base SigLIP model...")
+        # Fallback: load processor from base model online
+        processor = SiglipImageProcessor.from_pretrained("google/siglip-base-patch16-224")
+        print("✅ Image processor loaded from base model!")
     
-    # Get labels - handle case where id2label might not exist
+    # Get labels from your model config
     if hasattr(model.config, 'id2label') and model.config.id2label:
         labels = model.config.id2label
+        print(f"✅ Found {len(labels)} labels in model config")
     else:
         # Create generic labels if none exist
-        num_labels = model.config.num_labels if hasattr(model.config, 'num_labels') else 1000
+        num_labels = model.config.num_labels if hasattr(model.config, 'num_labels') else 2
         labels = {i: f"class_{i}" for i in range(num_labels)}
+        print(f"✅ Created {len(labels)} generic labels")
     
-    print(f"Model loaded successfully. Number of classes: {len(labels)}")
+    print("🎉 Model setup complete!")
     
 except Exception as e:
-    print(f"=== ERROR loading model from {MODEL_PATH} ===")
-    print(f"Error: {e}")
-    print("\n=== Debugging Information ===")
-    print("All files in Space:")
-    
-    def list_all_files(directory=".", prefix=""):
-        """Recursively list all files"""
-        try:
-            items = sorted(os.listdir(directory))
-            for item in items:
-                item_path = os.path.join(directory, item)
-                if os.path.isfile(item_path):
-                    size = os.path.getsize(item_path)
-                    print(f"{prefix}📄 {item} ({size} bytes)")
-                elif os.path.isdir(item_path) and not item.startswith('.'):
-                    print(f"{prefix}📁 {item}/")
-                    if len(prefix) < 6:  # Limit recursion depth
-                        list_all_files(item_path, prefix + "  ")
-        except PermissionError:
-            print(f"{prefix}❌ Permission denied")
-        except Exception as ex:
-            print(f"{prefix}❌ Error: {ex}")
-    
-    list_all_files()
-    
-    print("\n=== Required Files for Model ===")
-    print("✅ config.json - Model configuration")
-    print("✅ pytorch_model.bin OR model.safetensors - Model weights") 
-    print("✅ preprocessor_config.json - Image processor config")
-    print("✅ tokenizer.json (if applicable) - Tokenizer")
-    
-    print("\n=== Solutions ===")
-    print("1. Make sure all model files are uploaded to your Space")
-    print("2. Check that files aren't corrupted during upload")
-    print("3. Try uploading to a 'model' subfolder")
-    print("4. Verify the model was saved correctly during training")
-    
+    print(f"❌ Error loading model: {e}")
+    print("\n=== Debug Information ===")
+    print(f"Files in model directory: {os.listdir(MODEL_PATH)}")
     raise
 
-# Classify meme and extract text
 def classify_meme(image: Image.Image):
     """
     Classify meme and extract text using OCR
@@ -162,60 +56,61 @@ def classify_meme(image: Image.Image):
         # OCR: extract text from image
         extracted_text = pytesseract.image_to_string(image)
         
-        # Process image with the model
+        # Process image for the model
         inputs = processor(images=image, return_tensors="pt")
         
-        # Move inputs to same device as model if needed
-        if torch.cuda.is_available() and next(model.parameters()).is_cuda:
-            inputs = {k: v.to('cuda') for k, v in inputs.items()}
-        
+        # Run inference
         with torch.no_grad():
             outputs = model(**inputs)
             probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
         
-        # Get top predictions
-        top_k = min(10, len(labels))  # Show top 10 or all if fewer
-        top_probs, top_indices = torch.topk(probs[0], top_k)
-        
+        # Get predictions
         predictions = {}
-        for i, (prob, idx) in enumerate(zip(top_probs, top_indices)):
-            label = labels.get(idx.item(), f"class_{idx.item()}")
-            predictions[label] = float(prob)
+        for i in range(len(labels)):
+            label = labels.get(i, f"class_{i}")
+            predictions[label] = float(probs[0][i])
+        
+        # Sort predictions by confidence
+        sorted_predictions = dict(sorted(predictions.items(), key=lambda x: x[1], reverse=True))
         
-        # Debug prints (these will show in the console/logs)
-        print("Extracted Text:", extracted_text.strip())
-        print("Top Predictions:", predictions)
+        # Debug prints
+        print("=== Classification Results ===")
+        print(f"Extracted Text: '{extracted_text.strip()}'")
+        print("Top 3 Predictions:")
+        for i, (label, prob) in enumerate(list(sorted_predictions.items())[:3]):
+            print(f"  {i+1}. {label}: {prob:.4f}")
         
-        return predictions, extracted_text.strip()
+        return sorted_predictions, extracted_text.strip()
         
     except Exception as e:
-        print(f"Error in classification: {e}")
-        return {"Error": 1.0}, f"Error processing image: {str(e)}"
+        error_msg = f"Error processing image: {str(e)}"
+        print(f"❌ {error_msg}")
+        return {"Error": 1.0}, error_msg
 
-# Gradio interface
+# Create Gradio interface
 demo = gr.Interface(
     fn=classify_meme,
     inputs=gr.Image(type="pil", label="Upload Meme Image"),
     outputs=[
         gr.Label(num_top_classes=5, label="Meme Classification"),
-        gr.Textbox(label="Extracted Text from OCR", lines=3)
+        gr.Textbox(label="Extracted Text", lines=3)
     ],
-    title="Meme Classifier with OCR",
+    title="🎭 Meme Classifier with OCR",
     description="""
     Upload a meme image to:
-    1. Classify its content using your trained SigLIP2_77 model
-    2. Extract text using OCR (Optical Character Recognition)
+    1. **Classify** its content using your trained SigLIP2_77 model
+    2. **Extract text** using OCR (Optical Character Recognition)
     
-    Note: Make sure all model files are properly uploaded to your Space.
+    Your model was trained on meme data and will predict the category/sentiment of the uploaded meme.
     """,
     examples=None,
     allow_flagging="never"
 )
 
 if __name__ == "__main__":
-    print("Starting Gradio interface...")
+    print("🚀 Starting Gradio interface...")
     demo.launch(
-        server_name="0.0.0.0",  # Allow external connections in HF Spaces
-        server_port=7860,       # Standard port for HF Spaces
-        share=False             # HF Spaces handles sharing
-    )
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
+    )