Upload 18 files

.gitattributes

@@ -1,35 +1,4 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+open-images-dataset-train0.tsv filter=lfs diff=lfs merge=lfs -text
+News_Category_Dataset_v3.json filter=lfs diff=lfs merge=lfs -text
+custom_image_model.pth filter=lfs diff=lfs merge=lfs -text
+image_model.pth filter=lfs diff=lfs merge=lfs -text

analyze_model.py

@@ -0,0 +1,36 @@
+# analyze_model.py
+import torch
+from torchvision import models, transforms
+from torch.utils.data import DataLoader
+from torchvision.datasets import ImageFolder
+
+model = models.mobilenet_v2(weights=models.MobileNet_V2_Weights.IMAGENET1K_V1)
+model.classifier[1] = torch.nn.Linear(1280, 18)  # 18 classes
+model.load_state_dict(torch.load("custom_image_model.pth"))
+model.eval()
+
+transform = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+dataset = ImageFolder(root="categorized_images", transform=transform)
+val_loader = DataLoader(dataset, batch_size=16, shuffle=False)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+correct = 0
+total = 0
+with torch.no_grad():
+    for images, labels in val_loader:
+        images, labels = images.to(device), labels.to(device)
+        outputs = model(images)
+        _, predicted = torch.max(outputs, 1)
+        total += labels.size(0)
+        correct += (predicted == labels).sum().item()
+
+accuracy = 100 * correct / total
+print(f"✅ Model Accuracy: {accuracy:.2f}% on {total} images")

app.py

@@ -0,0 +1,65 @@
+import streamlit as st
+import os
+import shutil
+from PIL import Image
+import torch
+import torchvision.transforms as transforms
+from torchvision import models
+
+# Set up dataset path
+DATASET_PATH = "categorized_images"
+os.makedirs(DATASET_PATH, exist_ok=True)
+
+# Load class names dynamically from dataset folder
+class_names = sorted(os.listdir(DATASET_PATH))  # Get categories from folder names
+num_classes = len(class_names)
+
+# Load the trained model
+model = models.mobilenet_v2(weights=models.MobileNet_V2_Weights.IMAGENET1K_V1)
+model.classifier[1] = torch.nn.Linear(1280, num_classes)
+model.load_state_dict(torch.load("custom_image_model.pth", map_location=torch.device('cpu')))
+model.eval()
+
+# Define image transformation
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+def predict_and_save(image, filename):
+    """Predict category and save the image in the correct folder."""
+    image_tensor = transform(image).unsqueeze(0)
+    
+    with torch.no_grad():
+        output = model(image_tensor)
+        probabilities = torch.nn.functional.softmax(output, dim=1)
+        predicted_index = torch.argmax(probabilities, dim=1).item()
+    
+    predicted_category = class_names[predicted_index]
+    confidence = probabilities[0][predicted_index].item()
+    
+    # Ensure category folder exists
+    category_path = os.path.join(DATASET_PATH, predicted_category)
+    os.makedirs(category_path, exist_ok=True)
+    
+    # Save image in the correct category folder
+    image_save_path = os.path.join(category_path, filename)
+    image.save(image_save_path)
+    
+    return predicted_category, confidence, image_save_path
+
+# Streamlit UI
+st.title("📂 Smart Image Categorizer")
+st.write("Upload your images and let AI categorize them instantly!")
+
+uploaded_files = st.file_uploader("Upload images (single or multiple)", type=["png", "jpg", "jpeg"], accept_multiple_files=True)
+
+if uploaded_files:
+    for uploaded_file in uploaded_files:
+        image = Image.open(uploaded_file).convert("RGB")
+        category, confidence, saved_path = predict_and_save(image, uploaded_file.name)
+        
+        st.image(image, caption=f"{uploaded_file.name} → {category} ({confidence:.2%})", use_column_width=True)
+        st.success(f"✅ Categorized as: **{category}** (Confidence: {confidence:.2%})")
+        st.info(f"📂 Image saved to: {saved_path}")

categorizer.py

@@ -0,0 +1,58 @@
+import os
+import torch
+import torchvision.transforms as transforms
+from torchvision import models
+from PIL import Image
+
+# Load class names dynamically
+dataset_path = "categorized_images"
+class_names = sorted(os.listdir(dataset_path))  # Get categories from folder names
+num_classes = len(class_names)
+
+# Load trained model
+model = models.mobilenet_v2(weights=models.MobileNet_V2_Weights.IMAGENET1K_V1)
+model.classifier[1] = torch.nn.Linear(1280, num_classes)
+model.load_state_dict(torch.load("custom_image_model.pth", map_location=torch.device("cpu")))
+model.eval()
+
+# Define image transformations
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+def predict_category(image_path):
+    """Predicts the category of a single image."""
+    image = Image.open(image_path).convert("RGB")
+    image = transform(image).unsqueeze(0)
+    
+    with torch.no_grad():
+        output = model(image)
+        probabilities = torch.nn.functional.softmax(output, dim=1)
+        predicted_index = torch.argmax(probabilities, dim=1).item()
+    
+    return class_names[predicted_index]
+
+def categorize_images(image_folder="uncategorized_images", output_folder="categorized_images"):
+    """Categorizes all images in a folder."""
+    if not os.path.exists(image_folder):
+        print("❌ Image folder not found!")
+        return
+
+    for img_name in os.listdir(image_folder):
+        img_path = os.path.join(image_folder, img_name)
+        if not os.path.isfile(img_path):
+            continue
+        
+        category = predict_category(img_path)
+        category_folder = os.path.join(output_folder, category)
+        os.makedirs(category_folder, exist_ok=True)
+        
+        new_path = os.path.join(category_folder, img_name)
+        os.rename(img_path, new_path)
+        print(f"✅ Moved {img_name} to {category}/")
+
+if __name__ == "__main__":
+    categorize_images()
+    print("✅ Categorization complete!")

check_categories.py

@@ -0,0 +1,42 @@
+import os
+import json
+from collections import Counter
+from torchvision.datasets import ImageFolder
+
+# Paths
+dataset_path = "categorized_images"
+domain_config_path = "domain_config.json"
+
+# Load dataset using ImageFolder
+dataset = ImageFolder(root=dataset_path)
+
+# Count images in each class
+category_counts = Counter()
+for class_idx in dataset.targets:
+    category_counts[dataset.classes[class_idx]] += 1
+
+# Load domain_config.json
+with open(domain_config_path, "r") as f:
+    domain_config = json.load(f)
+
+# Print dataset classes and domain config keys
+print("\n✅ Dataset Classes from ImageFolder:", dataset.classes)
+print("\n✅ Categories in domain_config.json:", list(domain_config.keys()))
+
+# Check if classes match
+if set(dataset.classes) == set(domain_config.keys()):
+    print("\n✅ Class labels MATCH between dataset and domain_config.json!")
+else:
+    print("\n⚠️ WARNING: Mismatch between dataset classes and domain_config.json!")
+
+# Print category counts
+print("\n📊 Image Count Per Category:")
+for category, count in category_counts.items():
+    print(f"  - {category}: {count} images")
+
+# Check for empty categories
+empty_categories = [c for c in dataset.classes if category_counts[c] == 0]
+if empty_categories:
+    print("\n⚠️ WARNING: Some categories have 0 images:", empty_categories)
+else:
+    print("\n✅ All categories have images!")

check_dataset.py

@@ -0,0 +1,47 @@
+import os
+import json
+from torchvision.datasets import ImageFolder
+from torchvision import transforms
+
+# Load domain configuration
+config_path = "domain_config.json"
+with open(config_path, "r") as f:
+    domain_config = json.load(f)
+
+# Extract category names from domain_config.json
+config_categories = list(domain_config.keys())
+
+# Path to categorized images folder
+dataset_path = "categorized_images"
+
+# Apply data augmentation
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),  
+    transforms.RandomHorizontalFlip(p=0.5),   # Flip images randomly
+    transforms.RandomRotation(degrees=15),   # Rotate images by up to 15 degrees
+    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),  # Adjust colors
+    transforms.RandomResizedCrop(size=(224, 224), scale=(0.8, 1.0)),  # Random crop
+    transforms.ToTensor()
+])
+
+# Load dataset with augmentation
+dataset = ImageFolder(root=dataset_path, transform=transform)
+
+# Extract dataset categories
+dataset_categories = dataset.classes
+
+# Check for inconsistencies
+print("\n✅ Dataset Classes from ImageFolder:", dataset_categories)
+print("\n✅ Categories in domain_config.json:", config_categories)
+
+if set(dataset_categories) != set(config_categories):
+    print("\n⚠️ WARNING: Mismatch between dataset classes and domain_config.json!")
+
+# Count images per category
+print("\n📊 Image Count Per Category:")
+for category, idx in dataset.class_to_idx.items():
+    category_path = os.path.join(dataset_path, category)
+    num_images = len(os.listdir(category_path)) if os.path.exists(category_path) else 0
+    print(f"  - {category}: {num_images} images")
+
+print("\n✅ All categories have images!")

custom_image_model.pth

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

domain_config.json

@@ -0,0 +1,92 @@
+{
+    "Technology": {
+        "descriptions": ["a screenshot of a computer interface", "software application", "tech product"],
+        "keywords": ["code", "programming", "tech", "software"],
+        "weight": 1.0
+    },
+    "Finance": {
+        "descriptions": ["financial dashboard", "banking interface", "stock market"],
+        "keywords": ["bank", "money", "finance", "trading"],
+        "weight": 1.0
+    },
+    "Education": {
+        "descriptions": ["online learning", "educational website", "study materials"],
+        "keywords": ["study", "education", "course", "learn"],
+        "weight": 1.0
+    },
+    "Travel": {
+        "descriptions": ["travel booking", "maps", "travel planning interface"],
+        "keywords": ["travel", "ticket", "destination", "booking"],
+        "weight": 1.0
+    },
+    "Entertainment": {
+        "descriptions": ["streaming platform", "game interface", "media website"],
+        "keywords": ["movie", "game", "netflix", "youtube"],
+        "weight": 1.0
+    },
+    "E-commerce": {
+        "descriptions": ["shopping website", "product page", "marketplace"],
+        "keywords": ["buy", "cart", "shop", "product"],
+        "weight": 1.0
+    },
+    "Social Media": {
+        "descriptions": ["social app", "messaging platform", "social network interface"],
+        "keywords": ["facebook", "chat", "twitter", "instagram"],
+        "weight": 1.0
+    },
+    "News": {
+        "descriptions": ["news article", "news website", "online newspaper"],
+        "keywords": ["headline", "report", "news", "article"],
+        "weight": 1.0
+    },
+    "Productivity": {
+        "descriptions": ["task management", "project tracking", "productivity tool"],
+        "keywords": ["todo", "project", "management", "task"],
+        "weight": 1.0
+    },
+    "Sports": {
+        "descriptions": ["sports news", "sports statistics dashboard", "match schedule application"],
+        "keywords": ["match", "league", "team", "tournament"],
+        "weight": 1.0
+    },
+    "Food & Dining": {
+        "descriptions": ["food delivery", "restaurant review website", "recipe platform"],
+        "keywords": ["restaurant", "recipe", "food", "meal"],
+        "weight": 1.0
+    },
+    "Automotive": {
+        "descriptions": ["car shopping", "vehicle rental", "automobile marketplace"],
+        "keywords": ["vehicle", "rental", "car", "auto"],
+        "weight": 1.0
+    },
+    "Government & Public Services": {
+        "descriptions": ["government website", "public service portal", "tax filing system"],
+        "keywords": ["public service", "tax", "government", "policy"],
+        "weight": 1.0
+    },
+    "Nature": {
+        "descriptions": ["natural scenery", "wildlife photography", "landscape images"],
+        "keywords": ["forest", "mountains", "ocean", "nature"],
+        "weight": 1.0
+    },
+    "Quotes": {
+        "descriptions": ["motivational quotes", "inspirational sayings", "daily wisdom"],
+        "keywords": ["motivation", "inspiration", "quote", "wisdom"],
+        "weight": 1.0
+    },
+    "Resources": {
+        "descriptions": ["learning materials", "skill development resources", "reference guides"],
+        "keywords": ["skills", "reference", "guide", "tutorial"],
+        "weight": 1.0
+    },
+    "Ronaldo": {
+        "descriptions": ["Cristiano Ronaldo images", "football highlights of Ronaldo", "Ronaldo fan pages"],
+        "keywords": ["ronaldo", "football", "soccer", "cr7"],
+        "weight": 1.0
+    },
+    "Motivation": {
+        "descriptions": ["motivational content", "inspirational images", "quote visuals"],
+        "keywords": ["motivation", "inspire", "quote", "uplift"],
+        "weight": 1.0
+    }
+}

download_images.py

@@ -0,0 +1,40 @@
+import pandas as pd
+import requests
+import os
+
+file_path = "open-images-dataset-train0.tsv"
+
+# Read TSV file, skipping the first row
+df = pd.read_csv(file_path, sep="\t", engine="python", skiprows=1, names=["ImageURL", "Subset", "ImageID"])
+
+# Print first few rows to verify
+print("First few rows of the cleaned dataset:")
+print(df.head())
+
+# Create a fixed category folder (since 'Subset' contains numbers, not real categories)
+output_folder = "open_images_v7/dataset"
+os.makedirs(output_folder, exist_ok=True)
+
+# Limit downloads to the first 100 images
+max_images = 100
+
+for index, row in df.iterrows():
+    if index >= max_images:
+        break  # Stop downloading after 100 images
+
+    image_url = row["ImageURL"]
+    image_id = row["ImageID"]
+
+    # Ensure the image filename ends with ".jpg"
+    image_path = os.path.join(output_folder, f"{image_id}.jpg")
+
+    try:
+        response = requests.get(image_url, timeout=10)
+        if response.status_code == 200:
+            with open(image_path, "wb") as f:
+                f.write(response.content)
+            print(f"✅ Downloaded: {image_id}.jpg")
+        else:
+            print(f"❌ Failed: {image_id}")
+    except Exception as e:
+        print(f"❌ Error downloading {image_id}: {e}")

downloadimages.py

@@ -0,0 +1,33 @@
+import os
+import pandas as pd
+import requests
+from tqdm import tqdm
+
+# Load the CSV file
+csv_file = "insparation.csv"  # Make sure this is the correct file name
+df = pd.read_csv(csv_file)
+print("Column Names in CSV:", df.columns.tolist())
+
+# Ensure the column name matches your file
+url_column = "Image-link"  # Change this if the column name is different
+
+# Destination folder
+save_folder = "Motivation"
+os.makedirs(save_folder, exist_ok=True)
+
+# Set limit to 80 images
+num_images = min(80, len(df))  # If there are less than 80 URLs, take all available
+
+# Download images
+for idx, url in tqdm(enumerate(df[url_column][:num_images]), total=num_images):
+    try:
+        response = requests.get(url, stream=True)
+        if response.status_code == 200:
+            image_path = os.path.join(save_folder, f"motivation_{idx+1}.jpg")
+            with open(image_path, "wb") as file:
+                for chunk in response.iter_content(1024):
+                    file.write(chunk)
+    except Exception as e:
+        print(f"Failed to download {url}: {e}")
+
+print(f"Downloaded {num_images} images to {save_folder}")

huggingfacedownload.py.py

@@ -0,0 +1,7 @@
+from huggingface_hub import snapshot_download
+
+# Dataset repository name (from the URL)
+repo_id = "YashJain/UI-Elements-Detection-Dataset"
+
+# Download entire dataset
+snapshot_download(repo_id, repo_type="dataset", local_dir="UI_Dataset")

image_model.pth

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

predict.py

@@ -0,0 +1,76 @@
+import torch
+import torchvision.transforms as transforms
+from torchvision import models
+from PIL import Image
+import os
+import shutil
+import sys
+
+# Load class names dynamically from dataset folder
+dataset_path = "categorized_images"
+class_names = sorted(os.listdir(dataset_path))  # Get categories from folder names
+num_classes = len(class_names)
+
+# Load trained model
+model = models.mobilenet_v2(weights=models.MobileNet_V2_Weights.IMAGENET1K_V1)
+model.classifier[1] = torch.nn.Linear(1280, num_classes)
+model.load_state_dict(torch.load("custom_image_model.pth", map_location=torch.device('cpu')))
+model.eval()
+
+# Image transformation
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+def predict_and_categorize(image_path, move=True):
+    """Predict category for an image and move it to the correct folder."""
+    try:
+        image = Image.open(image_path).convert("RGB")
+    except Exception as e:
+        print(f"⚠️ Error loading image: {e}")
+        return
+
+    image_tensor = transform(image).unsqueeze(0)
+
+    with torch.no_grad():
+        output = model(image_tensor)
+        probabilities = torch.nn.functional.softmax(output, dim=1)
+        predicted_index = torch.argmax(probabilities, dim=1).item()
+
+    predicted_category = class_names[predicted_index]
+    confidence = probabilities[0][predicted_index].item()
+
+    print(f"✅ {image_path} -> **Predicted Category:** {predicted_category} ({confidence:.2%} confidence)")
+
+    # Move image to categorized_images folder
+    if move:
+        category_folder = os.path.join("categorized_images", predicted_category)
+        os.makedirs(category_folder, exist_ok=True)
+        shutil.move(image_path, os.path.join(category_folder, os.path.basename(image_path)))
+        print(f"📂 Moved to: {category_folder}\n")
+
+def process_folder(folder_path):
+    """Process all images in a folder."""
+    if not os.path.exists(folder_path):
+        print(f"❌ Folder not found: {folder_path}")
+        return
+
+    for file in os.listdir(folder_path):
+        if file.lower().endswith((".png", ".jpg", ".jpeg")):
+            predict_and_categorize(os.path.join(folder_path, file))
+
+if __name__ == "__main__":
+    if len(sys.argv) > 1:
+        input_path = sys.argv[1]
+
+        if os.path.isdir(input_path):
+            print(f"\n📂 **Processing folder:** {input_path}\n")
+            process_folder(input_path)
+        elif os.path.isfile(input_path):
+            predict_and_categorize(input_path)
+        else:
+            print("❌ Invalid path. Please provide an image or folder.")
+    else:
+        print("⚠️ Please provide an image or folder path.")

readingfile.py

@@ -0,0 +1,36 @@
+import requests
+import os
+
+# List of image URLs (Replace this with actual image links)
+image_urls = [
+    "https://example.com/image1.jpg",  # Replace with real image URLs
+    "https://example.com/image2.jpg",
+    # Add more image URLs here
+]
+
+# Folder to save images
+save_folder = "categorized_images/News"
+os.makedirs(save_folder, exist_ok=True)
+
+# Function to download images
+def download_image(url, folder):
+    try:
+        response = requests.get(url, headers={"User-Agent": "Mozilla/5.0"}, timeout=10)
+
+        # Check if response is an image
+        if "image" in response.headers["Content-Type"]:
+            filename = os.path.join(folder, url.split("/")[-1])
+            with open(filename, "wb") as file:
+                file.write(response.content)
+            print(f"✅ Downloaded: {filename}")
+        else:
+            print(f"❌ Not an image: {url}")
+
+    except Exception as e:
+        print(f"⚠️ Error downloading {url}: {e}")
+
+# Download only the first 80 images
+for index, url in enumerate(image_urls[:80]):  # Limit to 80 images
+    download_image(url, save_folder)
+
+print("\n🎉 Done! Downloaded up to 80 images.")

removebadimage.py

@@ -0,0 +1,15 @@
+from PIL import Image
+import os
+
+folder_path = "categorized_images"  # Update with your dataset folder
+
+for root, _, files in os.walk(folder_path):
+    for file in files:
+        file_path = os.path.join(root, file)
+        try:
+            with Image.open(file_path) as img:
+                img.verify()  # Verify if it's a valid image
+        except Exception as e:
+            print(f"Corrupt image found: {file_path}, Error: {e}")
+            os.remove(file_path)  # Remove corrupt image
+            print(f"Deleted: {file_path}")

scrapping.py

@@ -0,0 +1,57 @@
+import os
+import time
+import requests
+from selenium import webdriver
+from selenium.webdriver.chrome.service import Service
+from selenium.webdriver.common.by import By
+from webdriver_manager.chrome import ChromeDriverManager
+from tqdm import tqdm  # Progress bar
+
+# Setup Chrome Driver
+options = webdriver.ChromeOptions()
+options.add_argument("--headless")  # Run in background
+options.add_argument("--disable-gpu")  # Prevents rendering issues
+driver = webdriver.Chrome(service=Service(ChromeDriverManager().install()), options=options)
+
+# Open Pexels search page
+search_url = "https://www.pexels.com/search/productivity/"
+driver.get(search_url)
+
+# Wait for images to load
+time.sleep(5)
+
+# Scroll down multiple times to load more images
+for _ in range(10):  
+    driver.execute_script("window.scrollBy(0, 2000);")
+    time.sleep(2)  # Wait for new images to load
+
+# Find all image elements
+images = driver.find_elements(By.TAG_NAME, "img")
+
+# Extract Image URLs
+image_urls = []
+for img in images:
+    url = img.get_attribute("src")
+    if url and "pexels.com" in url:  # Ensure it's a valid image link
+        image_urls.append(url)
+
+# Keep only the first 100 images
+image_urls = image_urls[:100]
+
+# Create folder if not exists
+save_folder = "Productivity"
+os.makedirs(save_folder, exist_ok=True)
+
+# Download and save images
+for idx, img_url in enumerate(tqdm(image_urls, desc="Downloading Images")):
+    try:
+        img_data = requests.get(img_url).content
+        with open(os.path.join(save_folder, f"image_{idx+1}.jpg"), "wb") as f:
+            f.write(img_data)
+    except Exception as e:
+        print(f"Error downloading image {idx+1}: {e}")
+
+# Close the browser
+driver.quit()
+
+print(f"✅ {len(image_urls)} images downloaded in the '{save_folder}' folder.")

streamlit_app.py

@@ -0,0 +1,42 @@
+import streamlit as st
+import torch
+import torchvision.transforms as transforms
+from torchvision import models
+from PIL import Image
+import json
+import os
+
+with open("domain_config.json", "r") as f:
+    domain_config = json.load(f)
+    class_names = list(domain_config.keys())
+
+num_classes = len(class_names)
+
+model = models.mobilenet_v2(weights=models.MobileNet_V2_Weights.IMAGENET1K_V1)
+model.classifier[1] = torch.nn.Linear(1280, num_classes)
+model.load_state_dict(torch.load("custom_image_model.pth", map_location=torch.device('cpu')))
+model.eval()
+
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+st.title("📂 AI-Powered Image Categorization")
+
+uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+if uploaded_file is not None:
+    image = Image.open(uploaded_file).convert("RGB")
+    st.image(image, caption="Uploaded Image", use_column_width=True)
+
+    if st.button("Categorize Image"):
+        image_tensor = transform(image).unsqueeze(0)
+        with torch.no_grad():
+            output = model(image_tensor)
+            probabilities = torch.nn.functional.softmax(output, dim=1)
+            predicted_index = torch.argmax(probabilities, dim=1).item()
+
+        predicted_category = class_names[predicted_index]
+        confidence = probabilities[0][predicted_index].item()
+        st.success(f"✅ **Predicted Category:** {predicted_category} ({confidence:.2%} confidence)")

train_model.py

@@ -0,0 +1,104 @@
+import torch
+import torchvision.transforms as transforms
+from torch.utils.data import DataLoader, random_split
+import torchvision.models as models
+import torch.nn as nn
+import torch.optim as optim
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+from torchvision.datasets import ImageFolder
+import os
+
+def main():
+    dataset_path = "categorized_images"
+    if not os.path.exists(dataset_path):
+        raise FileNotFoundError(f"❌ Dataset folder '{dataset_path}' not found!")
+
+    # Get class names dynamically from dataset folders
+    class_names = sorted(os.listdir(dataset_path))
+    num_classes = len(class_names)
+
+    # Data Augmentation & Normalization
+    train_transform = transforms.Compose([
+        transforms.RandomResizedCrop(224),
+        transforms.RandomHorizontalFlip(),
+        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+
+    val_transform = transforms.Compose([
+        transforms.Resize(256),
+        transforms.CenterCrop(224),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+
+    dataset = ImageFolder(root=dataset_path, transform=train_transform)
+    train_size = int(0.8 * len(dataset))
+    val_size = len(dataset) - train_size
+    train_dataset, val_dataset = random_split(dataset, [train_size, val_size])
+
+    train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True, num_workers=4, pin_memory=True)
+    val_loader = DataLoader(val_dataset, batch_size=16, shuffle=False, num_workers=4, pin_memory=True)
+
+    # Load Pretrained Model
+    model = models.mobilenet_v2(weights=models.MobileNet_V2_Weights.IMAGENET1K_V1)
+    
+    # Freeze all layers except the classifier
+    for param in model.features.parameters():
+        param.requires_grad = False
+
+    # Update the classifier for our dataset
+    model.classifier[1] = nn.Linear(1280, num_classes)
+    
+    # Unfreeze last 3 layers to fine-tune
+    for param in model.features[-3:].parameters():
+        param.requires_grad = True
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.Adam(model.parameters(), lr=0.0001)
+    scheduler = ReduceLROnPlateau(optimizer, 'min', patience=3, factor=0.1)
+
+    best_val_loss = float('inf')
+    for epoch in range(30):
+        model.train()
+        train_loss = 0.0
+        for images, labels in train_loader:
+            images, labels = images.to(device), labels.to(device)
+            optimizer.zero_grad()
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+            train_loss += loss.item()
+        avg_train_loss = train_loss / len(train_loader)
+
+        model.eval()
+        val_loss, correct, total = 0.0, 0, 0
+        with torch.no_grad():
+            for images, labels in val_loader:
+                images, labels = images.to(device), labels.to(device)
+                outputs = model(images)
+                loss = criterion(outputs, labels)
+                val_loss += loss.item()
+                _, predicted = torch.max(outputs, 1)
+                total += labels.size(0)
+                correct += (predicted == labels).sum().item()
+        avg_val_loss = val_loss / len(val_loader)
+        val_accuracy = 100 * correct / total
+
+        print(f"📢 Epoch [{epoch+1}/30] → Train Loss: {avg_train_loss:.4f} | Val Loss: {avg_val_loss:.4f} | Val Accuracy: {val_accuracy:.2f}%")
+        scheduler.step(avg_val_loss)
+
+        if avg_val_loss < best_val_loss:
+            best_val_loss = avg_val_loss
+            torch.save(model.state_dict(), "custom_image_model.pth")
+            print("✅ Best model saved!")
+
+    print("🎉 Training Complete!")
+
+if __name__ == '__main__':
+    main()