Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

How to train gray pictures? #121

Open
WathomeBo opened this issue Dec 24, 2023 · 2 comments
Open

How to train gray pictures? #121

WathomeBo opened this issue Dec 24, 2023 · 2 comments

Comments

@WathomeBo
Copy link

When the image_dataset loads pictures, it will automatically convert the pictrue to rgb. However, i just want to keep the original type

@tobi-ore
Copy link

To train the model for 1 channel, you need to go to the image_datasets.py file and comment out line 97: arr = np.array(pil_image.convert("RGB")). This line forces the input channel to be RGB.

Then add this 2 lines after the line you commented:
arr = np.array(pil_image)
arr = arr.reshape((self.resolution, self.resolution, 1))

@Saumya-Gupta-26
Copy link

I made the following changes, and it worked:

In image_datasets.py,

# arr = np.array(pil_image.convert("RGB")) # commenting this line
arr = np.array(pil_image) # add this line
arr = arr.reshape((arr.shape[0], arr.shape[1], 1)) # add this line

In script_util.py, change the in_channels and out_channels variables being sent to UNetModel. Specifically:

return UNetModel(
    in_channels=1,
    model_channels=num_channels,
    out_channels=(1 if not learn_sigma else 2),

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

3 participants