Upload minimal files only

README.md

@@ -1,3 +1,79 @@
----
-license: apache-2.0
----
+# 🎨 Simple Colorizer - Image Colorization Model
+
+This repository contains a PyTorch-trained U-Net model that automatically colorizes grayscale images.
+
+---
+
+## 📂 Repository Contents
+
+- `best_colorization_model.pth`: Trained model weights
+- `model.py`: The `ImprovedUNet` architecture definition
+- `README.md`: This file
+
+---
+
+## 🚀 Usage Example
+
+### 1️⃣ Install Dependencies
+
+```python
+pip install -r requirements.txt
+```
+
+### 2️⃣ Load the Model
+```python
+import torch
+from model import ImprovedUNet
+```
+
+# Create the model instance
+```python
+model = ImprovedUNet()
+
+# Load the weights
+checkpoint = torch.load("best_colorization_model.pth", map_location="cpu")
+model.load_state_dict(checkpoint["model_state_dict"])
+model.eval()
+```
+### 3️⃣ Colorize an Image
+```python
+from PIL import Image
+import torchvision.transforms as T
+
+img = Image.open("path/to/grayscale_image.jpg").convert("L")
+transform = T.Compose([
+    T.Resize((256, 256)),
+    T.ToTensor(),
+    T.Normalize(mean=[0.5], std=[0.5])
+])
+
+input_tensor = transform(img).unsqueeze(0) 
+
+with torch.no_grad():
+    output = model(input_tensor)
+
+output_image = output.squeeze(0).permute(1, 2, 0).numpy()
+output_image = (output_image * 255).clip(0, 255).astype("uint8")
+
+Image.fromarray(output_image).save("colorized_output.png")
+```
+
+ℹ️ Training Information
+
+Architecture: Custom U-Net (ImprovedUNet)
+
+Input Size: 256x256 pixels
+
+Optimizer: Adam
+
+Loss Function: MSE
+
+Epochs: [Specify the number of epochs]
+
+📈 Results
+Here is an example of an image colorized by the model:
+![Colorized Example](test_result_PARIS.png)
+
+
+✨ Author
+This model was developed by Eric Houzelle.

colorizer-tiny.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:67e357b9c98f66b82ae79e5202110fae3d528b625c484e575463b9714256f74e
+size 372679415

model.py

@@ -0,0 +1,67 @@
+class ConvBlock(nn.Module):
+    def __init__(self, in_ch, out_ch, dropout=0.1):
+        super().__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_ch, out_ch, 3, padding=1),
+            nn.BatchNorm2d(out_ch),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_ch, out_ch, 3, padding=1),
+            nn.BatchNorm2d(out_ch),
+            nn.ReLU(inplace=True),
+            nn.Dropout2d(dropout)
+        )
+    def forward(self, x):
+        return self.conv(x)
+
+class ImprovedUNet(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.enc1 = ConvBlock(1, 64, dropout=0.1)
+        self.enc2 = ConvBlock(64, 128, dropout=0.1)
+        self.enc3 = ConvBlock(128, 256, dropout=0.2)
+        self.enc4 = ConvBlock(256, 512, dropout=0.2)
+        
+        self.pool = nn.MaxPool2d(2)
+        
+        self.bottleneck = ConvBlock(512, 1024, dropout=0.3)
+        
+        self.up4 = nn.ConvTranspose2d(1024, 512, 2, stride=2)
+        self.dec4 = ConvBlock(1024, 512, dropout=0.2)
+        
+        self.up3 = nn.ConvTranspose2d(512, 256, 2, stride=2)
+        self.dec3 = ConvBlock(512, 256, dropout=0.2)
+        
+        self.up2 = nn.ConvTranspose2d(256, 128, 2, stride=2)
+        self.dec2 = ConvBlock(256, 128, dropout=0.1)
+        
+        self.up1 = nn.ConvTranspose2d(128, 64, 2, stride=2)
+        self.dec1 = ConvBlock(128, 64, dropout=0.1)
+        
+        self.out_conv = nn.Conv2d(64, 3, 1)
+        self.out_act = nn.Tanh()
+    def forward(self, x):
+        e1 = self.enc1(x)
+        e2 = self.enc2(self.pool(e1))
+        e3 = self.enc3(self.pool(e2))
+        e4 = self.enc4(self.pool(e3))
+        
+        b = self.bottleneck(self.pool(e4))
+        
+        d4 = self.up4(b)
+        d4 = torch.cat([d4, e4], dim=1)
+        d4 = self.dec4(d4)
+        
+        d3 = self.up3(d4)
+        d3 = torch.cat([d3, e3], dim=1)
+        d3 = self.dec3(d3)
+        
+        d2 = self.up2(d3)
+        d2 = torch.cat([d2, e2], dim=1)
+        d2 = self.dec2(d2)
+        
+        d1 = self.up1(d2)
+        d1 = torch.cat([d1, e1], dim=1)
+        d1 = self.dec1(d1)
+        
+        out = self.out_conv(d1)
+        return self.out_act(out)