Update app.py

app.py

@@ -1,129 +1,26 @@
 import os
-import zipfile
 import torch
-from torch import nn, optim
-from torch.utils.data import DataLoader, Dataset
 from torchvision import transforms
 from PIL import Image
 from transformers import CLIPModel, CLIPProcessor
 import gradio as gr
 
-# Ensure PyTorch is installed
-try:
-    import torch
-except ModuleNotFoundError:
-    print("PyTorch is not installed. Installing now...")
-    os.system("pip install torch torchvision torchaudio")
-    import torch
-
-# Step 1: Unzip the dataset
-if not os.path.exists("data"):
-    os.makedirs("data")
-
-print("Extracting Data.zip...")
-with zipfile.ZipFile("Data.zip", 'r') as zip_ref:
-    zip_ref.extractall("data")
-print("Extraction complete.")
-
-# Step 2: Dynamically find the 'safe' and 'unsafe' folders
-def find_dataset_path(root_dir):
-    for root, dirs, files in os.walk(root_dir):
-        if 'safe' in dirs and 'unsafe' in dirs:
-            return root
-    return None
-
-# Look for 'safe' and 'unsafe' inside 'data/Data'
-dataset_path = find_dataset_path("data/Data")
-if dataset_path is None:
-    print("Debugging extracted structure:")
-    for root, dirs, files in os.walk("data"):
-        print(f"Root: {root}")
-        print(f"Directories: {dirs}")
-        print(f"Files: {files}")
-    raise FileNotFoundError("Expected 'safe' and 'unsafe' folders not found inside 'data/Data'. Please check the Data.zip structure.")
-print(f"Dataset path found: {dataset_path}")
-
-# Step 3: Define Custom Dataset Class
-class CustomImageDataset(Dataset):
-    def __init__(self, root_dir, transform=None):
-        self.root_dir = root_dir
-        self.transform = transform
-        self.image_paths = []
-        self.labels = []
-
-        for label, folder in enumerate(["safe", "unsafe"]):  # 0 = safe, 1 = unsafe
-            folder_path = os.path.join(root_dir, folder)
-            if not os.path.exists(folder_path):
-                raise FileNotFoundError(f"Folder '{folder}' not found in '{root_dir}'")
-            for filename in os.listdir(folder_path):
-                if filename.endswith((".jpg", ".jpeg", ".png")):  # Only load image files
-                    self.image_paths.append(os.path.join(folder_path, filename))
-                    self.labels.append(label)
-
-    def __len__(self):
-        return len(self.image_paths)
-
-    def __getitem__(self, idx):
-        image_path = self.image_paths[idx]
-        image = Image.open(image_path).convert("RGB")
-        label = self.labels[idx]
-        if self.transform:
-            image = self.transform(image)
-        return image, label
-
-# Step 4: Data Transformations
-transform = transforms.Compose([
-    transforms.Resize((224, 224)),  # Resize to 224x224 pixels
-    transforms.ToTensor(),         # Convert to tensor
-    transforms.Normalize((0.5,), (0.5,)),  # Normalize image values
-])
-
-# Step 5: Load the Dataset
-train_dataset = CustomImageDataset(dataset_path, transform=transform)
-train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
-
-# Debugging: Check the dataset
-print(f"Number of samples in the dataset: {len(train_dataset)}")
-if len(train_dataset) == 0:
-    raise ValueError("The dataset is empty. Please check if 'Data.zip' is correctly unzipped and contains 'safe' and 'unsafe' folders.")
-
-# Step 6: Load Pretrained CLIP Model
-model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
-processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
-
-# Add a Classification Layer
-model.classifier = nn.Linear(model.visual_projection.out_features, 2)  # 2 classes: safe, unsafe
-
-# Define Optimizer and Loss Function
-optimizer = optim.Adam(model.classifier.parameters(), lr=1e-4)
-criterion = nn.CrossEntropyLoss()
-
-# Step 7: Fine-Tune the Model
-model.train()
-for epoch in range(3):  # Number of epochs
-    total_loss = 0
-    for images, labels in train_loader:
-        optimizer.zero_grad()
-        images = torch.stack([img.to(torch.float32) for img in images])  # Batch of images
-        outputs = model.get_image_features(pixel_values=images)
-        logits = model.classifier(outputs)
-        loss = criterion(logits, labels)
-        loss.backward()
-        optimizer.step()
-        total_loss += loss.item()
-    print(f"Epoch {epoch+1}, Loss: {total_loss / len(train_loader)}")
-
-# Save the Fine-Tuned Model
-model.save_pretrained("fine-tuned-model")
-processor.save_pretrained("fine-tuned-model")
-print("Model fine-tuned and saved successfully.")
-
-# Step 8: Define Gradio Inference Function
+# Step 1: Ensure Fine-Tuned Model is Available
+fine_tuned_model_path = "fine-tuned-model"
+
+if not os.path.exists(fine_tuned_model_path):
+    raise FileNotFoundError(
+        f"The fine-tuned model is missing. Ensure that the fine-tuned model files are available in the '{fine_tuned_model_path}' directory."
+    )
+
+# Step 2: Load Fine-Tuned Model
+print("Loading fine-tuned model...")
+model = CLIPModel.from_pretrained(fine_tuned_model_path)
+processor = CLIPProcessor.from_pretrained(fine_tuned_model_path)
+print("Fine-tuned model loaded successfully.")
+
+# Step 3: Define Gradio Inference Function
 def classify_image(image, class_names):
-    # Load Fine-Tuned Model
-    model = CLIPModel.from_pretrained("fine-tuned-model")
-    processor = CLIPProcessor.from_pretrained("fine-tuned-model")
-    
     # Split class names from comma-separated input
     labels = [label.strip() for label in class_names.split(",") if label.strip()]
     if not labels:
@@ -139,7 +36,7 @@ def classify_image(image, class_names):
     result = {label: probs[0][i].item() for i, label in enumerate(labels)}
     return dict(sorted(result.items(), key=lambda item: item[1], reverse=True))
 
-# Step 9: Set Up Gradio Interface
+# Step 4: Set Up Gradio Interface
 iface = gr.Interface(
     fn=classify_image,
     inputs=[
@@ -151,7 +48,7 @@ iface = gr.Interface(
     description="Classify images as 'safe' or 'unsafe' using a fine-tuned CLIP model.",
 )
 
-# Launch Gradio Interface
+# Step 5: Launch Gradio Interface
 if __name__ == "__main__":
     iface.launch()
 
@@ -166,3 +63,4 @@ if __name__ == "__main__":
 
 
 
+