Update app.py

app.py

@@ -6,130 +6,10 @@ from huggingface_hub import hf_hub_download
 import torch
 import json
 from omegaconf import OmegaConf
-from model import Generator
-
-
-class BaseGenerator(ABC, nn.Module):
-    def __init__(self, channels: int = 3):
-        super().__init__()
-        self.channels = channels
-
-    @abstractmethod
-    def forward(self, x: Tensor) -> Tensor:
-        pass
-
-
-class Generator(BaseGenerator):
-    def __init__(self, cfg: DictConfig):
-        super().__init__(cfg.channels)
-        self.cfg = cfg
-        self.model = self._construct_model()
-
-    def _construct_model(self):
-        initial_layer = nn.Sequential(
-            nn.Conv2d(
-                self.cfg.channels,
-                self.cfg.num_features,
-                kernel_size=7,
-                stride=1,
-                padding=3,
-                padding_mode="reflect",
-            ),
-            nn.ReLU(inplace=True),
-        )
-
-        down_blocks = nn.Sequential(
-            ConvBlock(
-                self.cfg.num_features,
-                self.cfg.num_features * 2,
-                kernel_size=3,
-                stride=2,
-                padding=1,
-            ),
-            ConvBlock(
-                self.cfg.num_features * 2,
-                self.cfg.num_features * 4,
-                kernel_size=3,
-                stride=2,
-                padding=1,
-            ),
-        )
-
-        residual_blocks = nn.Sequential(
-            *[
-                ResidualBlock(self.cfg.num_features * 4)
-                for _ in range(self.cfg.num_residuals)
-            ]
-        )
-
-        up_blocks = nn.Sequential(
-            ConvBlock(
-                self.cfg.num_features * 4,
-                self.cfg.num_features * 2,
-                down=False,
-                kernel_size=3,
-                stride=2,
-                padding=1,
-                output_padding=1,
-            ),
-            ConvBlock(
-                self.cfg.num_features * 2,
-                self.cfg.num_features,
-                down=False,
-                kernel_size=3,
-                stride=2,
-                padding=1,
-                output_padding=1,
-            ),
-        )
-
-        last_layer = nn.Conv2d(
-            self.cfg.num_features,
-            self.cfg.channels,
-            kernel_size=7,
-            stride=1,
-            padding=3,
-            padding_mode="reflect",
-        )
-
-        return nn.Sequential(
-            initial_layer, down_blocks, residual_blocks, up_blocks, last_layer
-        )
-
-    def forward(self, x: Tensor) -> Tensor:
-        return tanh(self.model(x))
-
-
-class ConvBlock(nn.Module):
-    def __init__(
-        self, in_channels, out_channels, down=True, use_activation=True, **kwargs
-    ):
-        super().__init__()
-        self.conv = nn.Sequential(
-            nn.Conv2d(in_channels, out_channels, padding_mode="reflect", **kwargs)
-            if down
-            else nn.ConvTranspose2d(in_channels, out_channels, **kwargs),
-            nn.InstanceNorm2d(out_channels),
-            nn.ReLU(inplace=True) if use_activation else nn.Identity(),
-        )
-
-    def forward(self, x: Tensor) -> Tensor:
-        return self.conv(x)
-
-
-class ResidualBlock(nn.Module):
-    def __init__(self, channels: int):
-        super().__init__()
-        self.block = nn.Sequential(
-            ConvBlock(channels, channels, kernel_size=3, padding=1),
-            ConvBlock(
-                channels, channels, use_activation=False, kernel_size=3, padding=1
-            ),
-        )
-
-    def forward(self, x: Tensor) -> Tensor:
-        return x + self.block(x)
 
+import sys
+sys.path.append(os.path.dirname(model_path))
+from model import Generator
 
 repo_id = "Kiwinicki/sat2map-generator"
 generator_path = hf_hub_download(repo_id=repo_id, filename="generator.pth")
@@ -145,10 +25,23 @@ generator = Generator(cfg)
 generator.load_state_dict(torch.load(generator_path))
 generator.eval()
 
+from PIL import Image
+import torchvision.transforms as transforms
+
 
+transform = transforms.Compose([
+    transforms.Resize((256, 256)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+])
 
-def greet(iamge):
-    return image
+def process_image(image):
+    image_tensor = transform(image).unsqueeze(0)
+    with torch.no_grad():
+        output_tensor = generator(image_tensor)
+    output_image = output_tensor.squeeze(0)
+    output_image = transforms.ToPILImage()(output_image)
+    return output_image
 
-iface = gr.Interface(fn=greet, inputs="image", outputs="image")
-iface.launch()
+iface = gr.Interface(fn=process_image, inputs="image", outputs="image", title="Image Generator")
+iface.launch()