Spaces:

adarshcod30
/

drdetect-dr-screening

Running on Zero

drdetect-dr-screening / src /drdetect /fusion /model.py

Deploy DR screening Gradio demo (app.py fetches weights from the model repo)

b379113 verified about 1 month ago

945 Bytes

"""The fusion head itself: concat(CNN embedding, lesion features) -> ordinal head.
A small MLP, not another CNN -- the input is already a fixed-length vector
(the grading backbone's pooled embedding plus a handful of scalar lesion
features), not an image needing spatial processing.
"""

from __future__ import annotations

import torch
import torch.nn as nn

__all__ = ["FusionHead"]


class FusionHead(nn.Module):
    def __init__(
        self,
        *,
        embedding_dim: int,
        lesion_dim: int,
        hidden_dim: int = 128,
        num_outputs: int = 4,
        dropout: float = 0.3,
    ):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(embedding_dim + lesion_dim, hidden_dim),
            nn.ReLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, num_outputs),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)

	"""The fusion head itself: concat(CNN embedding, lesion features) -> ordinal head.

	A small MLP, not another CNN -- the input is already a fixed-length vector
	(the grading backbone's pooled embedding plus a handful of scalar lesion
	features), not an image needing spatial processing.
	"""

	from __future__ import annotations

	import torch
	import torch.nn as nn

	__all__ = ["FusionHead"]


	class FusionHead(nn.Module):
	def __init__(
	self,
	*,
	embedding_dim: int,
	lesion_dim: int,
	hidden_dim: int = 128,
	num_outputs: int = 4,
	dropout: float = 0.3,
	):
	super().__init__()
	self.net = nn.Sequential(
	nn.Linear(embedding_dim + lesion_dim, hidden_dim),
	nn.ReLU(inplace=True),
	nn.Dropout(dropout),
	nn.Linear(hidden_dim, num_outputs),
	)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	return self.net(x)