Update README.md

README.md

@@ -1,3 +1,11 @@
+---
+datasets:
+- phiyodr/coco2017
+language:
+- en
+library_name: fastai
+pipeline_tag: image-to-image
+---
 Model Card: GAN Colorization Model
 Model Description
 This GAN-based model performs image colorization, transforming grayscale images into color images. It leverages a generator network to predict the color channels and a discriminator network to improve the colorization quality through adversarial training.
@@ -22,85 +30,6 @@ Loss Functions:
 GAN Loss: Binary Cross-Entropy Loss with Logits
 L1 Loss: L1 Loss for pixel-wise comparison between generated and real color channels
 
-
-Model Inference
-
-torch
-torchvision
-fastai
-skimage
-matplotlib
-PIL
-numpy
-
-Inference Code
-
-from huggingface_hub import hf_hub_download
-import torch
-from PIL import Image
-from torchvision import transforms
-from skimage.color import rgb2lab, lab2rgb
-import numpy as np
-import matplotlib.pyplot as plt
-
-# Download the model from Hugging Face Hub
-repo_id = "Hammad712/GAN-Colorization-Model"
-model_filename = "generator.pt"
-model_path = hf_hub_download(repo_id=repo_id, filename=model_filename)
-
-# Define the generator model (same architecture as used during training)
-from fastai.vision.learner import create_body
-from torchvision.models import resnet34
-from fastai.vision.models.unet import DynamicUnet
-
-def build_generator(n_input=1, n_output=2, size=256):
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    backbone = create_body(resnet34(), pretrained=True, n_in=n_input, cut=-2)
-    G_net = DynamicUnet(backbone, n_output, (size, size)).to(device)
-    return G_net
-
-# Initialize and load the model
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-G_net = build_generator(n_input=1, n_output=2, size=256)
-G_net.load_state_dict(torch.load(model_path, map_location=device))
-G_net.eval()
-
-# Preprocessing function
-def preprocess_image(img_path):
-    img = Image.open(img_path).convert("RGB")
-    img = transforms.Resize((256, 256), Image.BICUBIC)(img)
-    img = np.array(img)
-    img_to_lab = rgb2lab(img).astype("float32")
-    img_to_lab = transforms.ToTensor()(img_to_lab)
-    L = img_to_lab[[0], ...] / 50. - 1.
-    return L.unsqueeze(0).to(device)
-
-# Inference function
-def colorize_image(img_path, model):
-    L = preprocess_image(img_path)
-    with torch.no_grad():
-        ab = model(L)
-    L = (L + 1.) * 50.
-    ab = ab * 110.
-    Lab = torch.cat([L, ab], dim=1).permute(0, 2, 3, 1).cpu().numpy()
-    rgb_imgs = []
-    for img in Lab:
-        img_rgb = lab2rgb(img)
-        rgb_imgs.append(img_rgb)
-    return np.stack(rgb_imgs, axis=0)
-
-# Example image path
-img_path = "/path/to/your/image.jpg"  # Replace with your image path
-
-# Perform inference
-colorized_images = colorize_image(img_path, G_net)
-
-# Display the result
-plt.imshow(colorized_images[0])
-plt.axis("off")
-plt.show()
-
-
 Usage
 To use the model for image colorization, ensure that the dependencies are installed and run the inference code provided. You will need to replace the image path with your own image for colorization.